Class: Whoosh::Auth::TokenTracker

Inherits:
Object
  • Object
show all
Defined in:
lib/whoosh/auth/token_tracker.rb

Instance Method Summary collapse

Constructor Details

#initializeTokenTracker

Returns a new instance of TokenTracker.



6
7
8
9
10
# File 'lib/whoosh/auth/token_tracker.rb', line 6

def initialize
  @usage = {}
  @callbacks = []
  @mutex = Mutex.new
end

Instance Method Details

#on_usage(&block) ⇒ Object



12
13
14
# File 'lib/whoosh/auth/token_tracker.rb', line 12

def on_usage(&block)
  @callbacks << block
end

#record(key:, endpoint:, tokens:) ⇒ Object



16
17
18
19
20
21
22
23
24
25
# File 'lib/whoosh/auth/token_tracker.rb', line 16

def record(key:, endpoint:, tokens:)
  total = tokens.values.sum
  @mutex.synchronize do
    @usage[key] ||= { total_tokens: 0, endpoints: {} }
    @usage[key][:total_tokens] += total
    @usage[key][:endpoints][endpoint] ||= 0
    @usage[key][:endpoints][endpoint] += total
  end
  @callbacks.each { |cb| cb.call(key, endpoint, tokens) }
end

#reset(key) ⇒ Object



35
36
37
# File 'lib/whoosh/auth/token_tracker.rb', line 35

def reset(key)
  @mutex.synchronize { @usage.delete(key) }
end

#usage_for(key) ⇒ Object



27
28
29
30
31
32
33
# File 'lib/whoosh/auth/token_tracker.rb', line 27

def usage_for(key)
  @mutex.synchronize do
    data = @usage[key]
    return { total_tokens: 0, endpoints: {} } unless data
    { total_tokens: data[:total_tokens], endpoints: data[:endpoints].dup }
  end
end