| /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| #ifndef TENSORFLOW_CORE_PLATFORM_CLOUD_GCS_THROTTLE_H_ |
| #define TENSORFLOW_CORE_PLATFORM_CLOUD_GCS_THROTTLE_H_ |
| |
| #include "tensorflow/core/platform/env.h" |
| |
| namespace tensorflow { |
| |
| /** |
| * GcsThrottleConfig is used to configure the GcsThrottle. |
| */ |
| struct GcsThrottleConfig { |
| /** |
| * enabled is true if GcsThrottle should throttle requests, false otherwise. |
| */ |
| bool enabled = false; |
| |
| /** |
| * token_rate is the number of tokens accrued every second that can be used |
| * for making requests to the GCS service. |
| */ |
| int64 token_rate = 100000; // Approximately 800 MBits/second bandwidth-only. |
| |
| /** |
| * bucket_size is the maximum number of available tokens the GcsThrottle can |
| * accrue. |
| */ |
| int64 bucket_size = 10000000; // 10 million tokens total |
| |
| /** |
| * tokens_per_request determines the number of tokens consumed for every |
| * request. |
| * |
| * Note: tokens are also consumed in proportion to the response size. |
| */ |
| int64 tokens_per_request = 100; |
| |
| /** |
| * initial_tokens determines how many tokens should be available immediately |
| * after the GcsThrottle is constructed. |
| */ |
| int64 initial_tokens = 0; |
| }; |
| |
| /** |
| * GcsThrottle is used to ensure fair use of the available GCS capacity. |
| * |
| * GcsThrottle operates around a concept of tokens. Tokens are consumed when |
| * making requests to the GCS service. Tokens are consumed both based on the |
| * number of requests made, as well as the bandwidth consumed (response sizes). |
| * |
| * GcsThrottle is thread safe and can be used from multiple threads. |
| */ |
| class GcsThrottle { |
| public: |
| /** |
| * Constructs a GcsThrottle. |
| */ |
| explicit GcsThrottle(EnvTime* env_time = EnvTime::Default()); |
| |
| /** |
| * AdmitRequest updates the GcsThrottle to record a request will be made. |
| * |
| * AdmitRequest should be called before any request is made. AdmitRequest |
| * returns false if the request should be denied. If AdmitRequest |
| * returns false, no tokens are consumed. If true is returned, the configured |
| * number of tokens are consumed. |
| */ |
| bool AdmitRequest(); |
| |
| /** |
| * RecordResponse updates the GcsThrottle to record a request has been made. |
| * |
| * RecordResponse should be called after the response has been received. |
| * RecordResponse will update the internal state based on the number of bytes |
| * in the response. |
| * |
| * Note: we split up the request and the response in this fashion in order to |
| * avoid penalizing consumers who are using large readahead buffers at higher |
| * layers of the I/O stack. |
| */ |
| void RecordResponse(size_t num_bytes); |
| |
| /** |
| * SetConfig sets the configuration for GcsThrottle and re-initializes state. |
| * |
| * After calling this, the token pool will be config.initial_tokens. |
| */ |
| void SetConfig(GcsThrottleConfig config); |
| |
| /** |
| * available_tokens gives a snapshot of how many tokens are available. |
| * |
| * The returned value should not be used to make admission decisions. The |
| * purpose of this function is to make available to monitoring or other |
| * instrumentation the number of available tokens in the pool. |
| */ |
| inline int64 available_tokens() LOCKS_EXCLUDED(mu_) { |
| mutex_lock l(mu_); |
| UpdateState(); |
| return available_tokens_; |
| } |
| |
| /** |
| * is_enabled determines if the throttle is enabled. |
| * |
| * If !is_enabled(), AdmitRequest() will always return true. To enable the |
| * throttle, call SetConfig passing in a configuration that has enabled set to |
| * true. |
| */ |
| bool is_enabled() LOCKS_EXCLUDED(mu_) { |
| mutex_lock l(mu_); |
| return config_.enabled; |
| } |
| |
| private: |
| /** |
| * UpdateState updates the available_tokens_ and last_updated_secs_ variables. |
| * |
| * UpdateState should be called in order to mark the passage of time, and |
| * therefore add tokens to the available_tokens_ pool. |
| */ |
| void UpdateState() EXCLUSIVE_LOCKS_REQUIRED(mu_); |
| |
| inline uint64 request_bytes_to_tokens(size_t num_bytes) { |
| return num_bytes >> 10; |
| } |
| |
| mutex mu_; |
| |
| /** |
| * last_updated_secs_ records the number of seconds since the Unix epoch that |
| * the internal state of the GcsThrottle was updated. This is important when |
| * determining the number of tokens to add to the available_tokens_ pool. |
| */ |
| uint64 last_updated_secs_ GUARDED_BY(mu_) = 0; |
| |
| /** |
| * available_tokens_ records how many tokens are available to be consumed. |
| * |
| * Note: it is possible for available_tokens_ to become negative. If a |
| * response comes back that consumes more than the available tokens, the count |
| * will go negative, and block future requests until we have available tokens. |
| */ |
| int64 available_tokens_ GUARDED_BY(mu_) = 0; |
| |
| EnvTime* const env_time_; |
| GcsThrottleConfig config_ GUARDED_BY(mu_); |
| }; |
| |
| } // namespace tensorflow |
| |
| #endif // TENSORFLOW_CORE_PLATFORM_CLOUD_GCS_THROTTLE_H_ |