Class: LightrateClient::TokenBucket

Inherits:
Object
  • Object
show all
Defined in:
lib/lightrate_client/types.rb

Overview

Token bucket for local token management

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(max_tokens, rule_id:, matcher:, http_method: nil, user_identifier:) ⇒ TokenBucket 



121
122
123
124
125
126
127
128
129
130
 
# File 'lib/lightrate_client/types.rb', line 121

def initialize(max_tokens, rule_id:, matcher:, http_method: nil, user_identifier:)
  @max_tokens = max_tokens
  @available_tokens = 0
  @rule_id = rule_id
  @matcher = matcher
  @http_method = http_method
  @last_accessed_at = Time.now
  @user_identifier = user_identifier
  @mutex = Mutex.new
end

Instance Attribute Details

#available_tokens ⇒ Object (readonly)

Returns the value of attribute available_tokens.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def available_tokens
  @available_tokens
end

#http_method ⇒ Object (readonly)

Returns the value of attribute http_method.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def http_method
  @http_method
end

#last_accessed_at ⇒ Object (readonly)

Returns the value of attribute last_accessed_at.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def last_accessed_at
  @last_accessed_at
end

#matcher ⇒ Object (readonly)

Returns the value of attribute matcher.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def matcher
  @matcher
end

#max_tokens ⇒ Object (readonly)

Returns the value of attribute max_tokens.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def max_tokens
  @max_tokens
end

#rule_id ⇒ Object (readonly)

Returns the value of attribute rule_id.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def rule_id
  @rule_id
end

#user_identifier ⇒ Object (readonly)

Returns the value of attribute user_identifier.



119
120
121
 
# File 'lib/lightrate_client/types.rb', line 119

def user_identifier
  @user_identifier
end

Instance Method Details

#check_and_consume_token ⇒ Array

Check tokens and consume atomically (caller must hold lock) This prevents race conditions between checking and consuming



225
226
227
228
229
230
231
232
233
234
235
236
 
# File 'lib/lightrate_client/types.rb', line 225

def check_and_consume_token
  synchronize do
    touch
    has_tokens = @available_tokens > 0
    if has_tokens
      @available_tokens -= 1
      true
    else
      false
    end
  end
end

#consume_token ⇒ Boolean

Consume one token from the bucket (caller must hold lock)



140
141
142
143
144
145
 
# File 'lib/lightrate_client/types.rb', line 140

def consume_token
  return false if @available_tokens <= 0
  
  @available_tokens -= 1
  true
end

#consume_tokens(count) ⇒ Integer

Consume multiple tokens from the bucket (caller must hold lock)



150
151
152
153
154
155
156
 
# File 'lib/lightrate_client/types.rb', line 150

def consume_tokens(count)
  return 0 if count <= 0 || @available_tokens <= 0
  
  tokens_to_consume = [count, @available_tokens].min
  @available_tokens -= tokens_to_consume
  tokens_to_consume
end

#expired? ⇒ Boolean

Check if bucket has expired (not accessed in 60 seconds)



213
214
215
 
# File 'lib/lightrate_client/types.rb', line 213

def expired?
  Time.now - @last_accessed_at > 60
end

#has_tokens? ⇒ Boolean

Check if tokens are available locally (caller must hold lock)



134
135
136
 
# File 'lib/lightrate_client/types.rb', line 134

def has_tokens?
  @available_tokens > 0
end

#matches?(operation, path, http_method) ⇒ Boolean

Check if this bucket matches the given request



183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
 
# File 'lib/lightrate_client/types.rb', line 183

def matches?(operation, path, http_method)
  return false if expired?
  return false unless @matcher
  
  begin
    matcher_regex = Regexp.new(@matcher)
    
    # For operation-based requests, match against operation
    if operation
      return matcher_regex.match?(operation) && @http_method.nil?
    end
    
    # For path-based requests, match against path and HTTP method
    if path
      return matcher_regex.match?(path) && @http_method == http_method
    end
    
    false
  rescue RegexpError
    # If matcher is not a valid regex, fall back to exact match
    if operation
      return @matcher == operation && @http_method.nil?
    elsif path
      return @matcher == path && @http_method == http_method
    end
    false
  end
end

#refill(tokens_to_add) ⇒ Integer

Refill the bucket with tokens from the server (caller must hold lock)



161
162
163
164
165
166
 
# File 'lib/lightrate_client/types.rb', line 161

def refill(tokens_to_add)
  touch
  tokens_to_add = [tokens_to_add, @max_tokens - @available_tokens].min
  @available_tokens += tokens_to_add
  tokens_to_add
end

#reset ⇒ Object

Reset bucket to empty state (caller must hold lock)



178
179
180
 
# File 'lib/lightrate_client/types.rb', line 178

def reset
  @available_tokens = 0
end

#status ⇒ Hash

Get current bucket status (caller must hold lock)



170
171
172
173
174
175
 
# File 'lib/lightrate_client/types.rb', line 170

def status
  {
    tokens_remaining: @available_tokens,
    max_tokens: @max_tokens
  }
end

#synchronize { ... } ⇒ Object

Synchronize access to this bucket for thread-safe operations

Yields:

  • Block to execute under bucket lock



240
241
242
 
# File 'lib/lightrate_client/types.rb', line 240

def synchronize(&block)
  @mutex.synchronize(&block)
end

#touch ⇒ Object

Update last accessed time



218
219
220
 
# File 'lib/lightrate_client/types.rb', line 218

def touch
  @last_accessed_at = Time.now
end