1 /* Copyright 2017 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_PLATFORM_CLOUD_GCS_THROTTLE_H_ 17 #define TENSORFLOW_CORE_PLATFORM_CLOUD_GCS_THROTTLE_H_ 18 19 #include "tensorflow/core/platform/env.h" 20 21 namespace tensorflow { 22 23 /** 24 * GcsThrottleConfig is used to configure the GcsThrottle. 25 */ 26 struct GcsThrottleConfig { 27 /** 28 * enabled is true if GcsThrottle should throttle requests, false otherwise. 29 */ 30 bool enabled = false; 31 32 /** 33 * token_rate is the number of tokens accrued every second that can be used 34 * for making requests to the GCS service. 35 */ 36 int64 token_rate = 100000; // Approximately 800 MBits/second bandwidth-only. 37 38 /** 39 * bucket_size is the maximum number of available tokens the GcsThrottle can 40 * accrue. 41 */ 42 int64 bucket_size = 10000000; // 10 million tokens total 43 44 /** 45 * tokens_per_request determines the number of tokens consumed for every 46 * request. 47 * 48 * Note: tokens are also consumed in proportion to the response size. 49 */ 50 int64 tokens_per_request = 100; 51 52 /** 53 * initial_tokens determines how many tokens should be available immediately 54 * after the GcsThrottle is constructed. 55 */ 56 int64 initial_tokens = 0; 57 }; 58 59 /** 60 * GcsThrottle is used to ensure fair use of the available GCS capacity. 61 * 62 * GcsThrottle operates around a concept of tokens. Tokens are consumed when 63 * making requests to the GCS service. Tokens are consumed both based on the 64 * number of requests made, as well as the bandwidth consumed (response sizes). 65 * 66 * GcsThrottle is thread safe and can be used from multiple threads. 67 */ 68 class GcsThrottle { 69 public: 70 /** 71 * Constructs a GcsThrottle. 72 */ 73 explicit GcsThrottle(EnvTime* env_time = nullptr); 74 75 /** 76 * AdmitRequest updates the GcsThrottle to record a request will be made. 77 * 78 * AdmitRequest should be called before any request is made. AdmitRequest 79 * returns false if the request should be denied. If AdmitRequest 80 * returns false, no tokens are consumed. If true is returned, the configured 81 * number of tokens are consumed. 82 */ 83 bool AdmitRequest(); 84 85 /** 86 * RecordResponse updates the GcsThrottle to record a request has been made. 87 * 88 * RecordResponse should be called after the response has been received. 89 * RecordResponse will update the internal state based on the number of bytes 90 * in the response. 91 * 92 * Note: we split up the request and the response in this fashion in order to 93 * avoid penalizing consumers who are using large readahead buffers at higher 94 * layers of the I/O stack. 95 */ 96 void RecordResponse(size_t num_bytes); 97 98 /** 99 * SetConfig sets the configuration for GcsThrottle and re-initializes state. 100 * 101 * After calling this, the token pool will be config.initial_tokens. 102 */ 103 void SetConfig(GcsThrottleConfig config); 104 105 /** 106 * available_tokens gives a snapshot of how many tokens are available. 107 * 108 * The returned value should not be used to make admission decisions. The 109 * purpose of this function is to make available to monitoring or other 110 * instrumentation the number of available tokens in the pool. 111 */ available_tokens()112 inline int64 available_tokens() LOCKS_EXCLUDED(mu_) { 113 mutex_lock l(mu_); 114 UpdateState(); 115 return available_tokens_; 116 } 117 118 /** 119 * is_enabled determines if the throttle is enabled. 120 * 121 * If !is_enabled(), AdmitRequest() will always return true. To enable the 122 * throttle, call SetConfig passing in a configuration that has enabled set to 123 * true. 124 */ is_enabled()125 bool is_enabled() LOCKS_EXCLUDED(mu_) { 126 mutex_lock l(mu_); 127 return config_.enabled; 128 } 129 130 private: 131 /** 132 * UpdateState updates the available_tokens_ and last_updated_secs_ variables. 133 * 134 * UpdateState should be called in order to mark the passage of time, and 135 * therefore add tokens to the available_tokens_ pool. 136 */ 137 void UpdateState() EXCLUSIVE_LOCKS_REQUIRED(mu_); 138 request_bytes_to_tokens(size_t num_bytes)139 inline uint64 request_bytes_to_tokens(size_t num_bytes) { 140 return num_bytes >> 10; 141 } 142 143 mutex mu_; 144 145 /** 146 * last_updated_secs_ records the number of seconds since the Unix epoch that 147 * the internal state of the GcsThrottle was updated. This is important when 148 * determining the number of tokens to add to the available_tokens_ pool. 149 */ 150 uint64 last_updated_secs_ GUARDED_BY(mu_) = 0; 151 152 /** 153 * available_tokens_ records how many tokens are available to be consumed. 154 * 155 * Note: it is possible for available_tokens_ to become negative. If a 156 * response comes back that consumes more than the available tokens, the count 157 * will go negative, and block future requests until we have available tokens. 158 */ 159 int64 available_tokens_ GUARDED_BY(mu_) = 0; 160 161 EnvTime* const env_time_; 162 GcsThrottleConfig config_ GUARDED_BY(mu_); 163 }; 164 165 } // namespace tensorflow 166 167 #endif // TENSORFLOW_CORE_PLATFORM_CLOUD_GCS_THROTTLE_H_ 168