• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 = EnvTime::Default());
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