1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_CHANNEL_COMMON_H_ 17 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_CHANNEL_COMMON_H_ 18 19 #include <unordered_map> 20 #include <vector> 21 22 #include "absl/container/flat_hash_map.h" 23 #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" 24 #include "tensorflow/core/platform/logging.h" 25 #include "tensorflow/core/platform/mutex.h" 26 27 namespace tensorflow { 28 29 // GenericCachingChannelCache that caches results to FindWorkerChannel() calls. 30 // To use instantiate with the type of channel cache needed. 31 // GenericCachingChannelCache allows using multiple channels to communiate with 32 // same target to provide throughput gains. When multiple channels exist for 33 // the same target they are chosen in a simple round robin fashion on each call 34 // to FindWorkerChannel. 35 template <typename ChannelCacheT> 36 class GenericCachingChannelCache : public ChannelCacheT { 37 public: GenericCachingChannelCache(int num_channels_per_target)38 explicit GenericCachingChannelCache(int num_channels_per_target) 39 : num_channels_per_target_( 40 num_channels_per_target > 0 ? num_channels_per_target : 1) {} 41 ~GenericCachingChannelCache()42 ~GenericCachingChannelCache() override {} 43 FindWorkerChannel(const string & target)44 SharedGrpcChannelPtr FindWorkerChannel(const string& target) override { 45 { 46 mutex_lock l(mu_); 47 auto iter = channels_.find(target); 48 if (iter != channels_.end()) { 49 return GetNextChannelPtrAndUpdateState(iter->second); 50 } 51 } 52 ChannelState new_chan_state; 53 for (int indx = 0; indx < num_channels_per_target_; indx++) { 54 auto ch = FindChannelOnce(target); 55 if (!ch) return nullptr; 56 new_chan_state.channels.push_back(ch); 57 } 58 new_chan_state.last_used = num_channels_per_target_ - 1; 59 60 { 61 mutex_lock l(mu_); 62 typename absl::flat_hash_map<string, ChannelState>::iterator iter; 63 bool was_inserted; 64 std::tie(iter, was_inserted) = channels_.insert({target, new_chan_state}); 65 VLOG(2) << "Channel cache for target: " << target 66 << " Size: " << new_chan_state.channels.size() 67 << " insertion: " << was_inserted; 68 return GetNextChannelPtrAndUpdateState(iter->second); 69 } 70 } 71 72 protected: 73 // Find the ClientChannel for "target". Only called when no channel was 74 // found in the channels_ cache for "target". A non nullptr result will be 75 // cached in channels_. 76 virtual SharedGrpcChannelPtr FindChannelOnce(const string& target) = 0; 77 78 private: 79 struct ChannelState { 80 std::vector<SharedGrpcChannelPtr> channels; 81 int last_used; 82 }; 83 84 // Should be called with mu_ held. GetNextChannelPtrAndUpdateState(ChannelState & chan_state)85 SharedGrpcChannelPtr GetNextChannelPtrAndUpdateState( 86 ChannelState& chan_state) { 87 // Following statement is marked as Crash OK as this is an invariant of 88 // code flow in this class. 89 CHECK_EQ(chan_state.channels.size(), num_channels_per_target_); // Crash OK 90 chan_state.last_used = 91 (chan_state.last_used + 1) % num_channels_per_target_; 92 return chan_state.channels[chan_state.last_used]; 93 } 94 95 const int num_channels_per_target_; 96 // TODO(zhifengc): Eviction when the map becomes too big. 97 mutex mu_; 98 absl::flat_hash_map<string, ChannelState> channels_ TF_GUARDED_BY(mu_); 99 }; 100 101 } // namespace tensorflow 102 103 #endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_CHANNEL_COMMON_H_ 104