1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15 16#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RECENT_REQUEST_IDS_H_ 17#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RECENT_REQUEST_IDS_H_ 18 19#include <vector> 20 21#include "tensorflow/core/lib/core/status.h" 22#include "tensorflow/core/lib/gtl/flatset.h" 23#include "tensorflow/core/platform/mutex.h" 24#include "tensorflow/core/platform/thread_annotations.h" 25#include "tensorflow/core/platform/types.h" 26#include "tensorflow/core/protobuf/worker.pb.h" 27 28namespace tensorflow { 29 30// RecentRequestIds tracks recent 64-bit request_ids. When maximum capacity is 31// reached, the oldest request_id is evicted. Thread safe. 32// 33// Some RPCs like RecvTensor are unsafe to retry. For example, RecvTensor pairs 34// one sender and one receiver, and the receiver waits for the sender's tensor. 35// Retried RecvTensor requests are problematic, because the original RecvTensor 36// request may have consumed the sender's tensor, so a retried request might 37// block forever. RecentRequestIds identifies retried requests, so we can fail 38// them instead of blocking forever. 39// 40// Internally, recent request_ids are stored in two data structures: a set and a 41// circular buffer. The set is used for efficient lookups, and the circular 42// buffer tracks the oldest request_id. When the buffer is full, the new 43// request_id replaces the oldest request_id in the circular buffer, and the 44// oldest request_id is removed from the set. 45class RecentRequestIds { 46 public: 47 // num_tracked_request_ids should be much larger than the number of RPCs that 48 // can be received in a small time window. For example, we observed a peak RPC 49 // rate of ~700 RecvTensor RPC/s when training inception v3 on TPUs, so we 50 // currently set num_tracked_request_ids to 100,000 for RecvTensor. 51 RecentRequestIds(int num_tracked_request_ids); 52 53 // Returns OK iff request_id has not been seen in the last 54 // num_tracked_request_ids insertions. For backwards compatibility, this 55 // always returns OK for request_id 0. The method_name and the request's 56 // ShortDebugString are added to returned errors. 57 Status TrackUnique(int64 request_id, const string& method_name, 58 const protobuf::Message& request); 59 60 private: 61 mutex mu_; 62 // next_index_ indexes into circular_buffer_, and points to the next storage 63 // space to use. When the buffer is full, next_index_ points at the oldest 64 // request_id. 65 int next_index_ GUARDED_BY(mu_) = 0; 66 std::vector<int64> circular_buffer_ GUARDED_BY(mu_); 67 gtl::FlatSet<int64> set_ GUARDED_BY(mu_); 68}; 69 70} // namespace tensorflow 71 72#endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RECENT_REQUEST_IDS_H_ 73