1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
%% This Source Code Form is subject to the terms of the Mozilla Public
%% License, v. 2.0. If a copy of the MPL was not distributed with this
%% file, You can obtain one at https://mozilla.org/MPL/2.0/.
%%
%% Copyright (c) 2007-2020 VMware, Inc. or its affiliates. All rights reserved.
%%
-module(rabbit_auth_cache_ets_segmented_stateless).
-behaviour(gen_server).
-behaviour(rabbit_auth_cache).
-export([start_link/1,
get/1, put/3, delete/1]).
-export([gc/0]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
terminate/2, code_change/3]).
-define(SEGMENT_TABLE, rabbit_auth_cache_ets_segmented_stateless_segment_table).
-record(state, {gc_timer}).
start_link(SegmentSize) ->
gen_server:start_link({local, ?MODULE}, ?MODULE, [SegmentSize], []).
get(Key) ->
case get_from_segments(Key) of
[] -> {error, not_found};
[V|_] -> {ok, V}
end.
put(Key, Value, TTL) ->
Expiration = rabbit_auth_cache:expiration(TTL),
[{_, SegmentSize}] = ets:lookup(?SEGMENT_TABLE, segment_size),
Segment = segment(Expiration, SegmentSize),
Table = case ets:lookup(?SEGMENT_TABLE, Segment) of
[{Segment, T}] -> T;
[] -> add_segment(Segment)
end,
ets:insert(Table, {Key, {Expiration, Value}}),
ok.
delete(Key) ->
[ets:delete(Table, Key)
|| Table <- get_all_segment_tables()].
gc() ->
case whereis(?MODULE) of
undefined -> ok;
Pid -> Pid ! gc
end.
init([SegmentSize]) ->
ets:new(?SEGMENT_TABLE, [ordered_set, named_table, public]),
ets:insert(?SEGMENT_TABLE, {segment_size, SegmentSize}),
InitSegment = segment(rabbit_auth_cache:expiration(SegmentSize), SegmentSize),
do_add_segment(InitSegment),
{ok, GCTimer} = timer:send_interval(SegmentSize * 2, gc),
{ok, #state{gc_timer = GCTimer}}.
handle_call({add_segment, Segment}, _From, State) ->
%% Double check segment if it's already created
Table = do_add_segment(Segment),
{reply, Table, State}.
handle_cast(_, State = #state{}) ->
{noreply, State}.
handle_info(gc, State = #state{}) ->
Now = erlang:system_time(milli_seconds),
MatchSpec = [{{'$1', '$2'}, [{'<', '$1', {const, Now}}], ['$2']}],
Expired = ets:select(?SEGMENT_TABLE, MatchSpec),
[ets:delete(Table) || Table <- Expired],
{noreply, State};
handle_info(_Msg, State) ->
{noreply, State}.
code_change(_OldVsn, State, _Extra) ->
{ok, State}.
terminate(_Reason, State = #state{gc_timer = Timer}) ->
timer:cancel(Timer),
State.
segment(Expiration, SegmentSize) ->
Begin = ((Expiration div SegmentSize) * SegmentSize),
End = Begin + SegmentSize,
End.
add_segment(Segment) ->
gen_server:call(?MODULE, {add_segment, Segment}).
do_add_segment(Segment) ->
case ets:lookup(?SEGMENT_TABLE, Segment) of
[{Segment, Table}] -> Table;
[] -> Table = ets:new(segment, [set, public]),
ets:insert(?SEGMENT_TABLE, {Segment, Table}),
Table
end.
get_segment_tables() ->
Now = erlang:system_time(milli_seconds),
MatchSpec = [{{'$1', '$2'}, [{'>', '$1', {const, Now}}], ['$_']}],
[V || {K, V} <- ets:select(?SEGMENT_TABLE, MatchSpec), K =/= segment_size].
get_all_segment_tables() ->
[V || {K, V} <- ets:tab2list(?SEGMENT_TABLE), K =/= segment_size].
get_from_segments(Key) ->
Tables = get_segment_tables(),
lists:flatmap(
fun(undefined) -> [];
(T) ->
try ets:lookup(T, Key) of
[{Key, {Exp, Val}}] ->
case rabbit_auth_cache:expired(Exp) of
true -> [];
false -> [Val]
end;
[] -> []
% ETS table can be deleted concurrently.
catch
error:badarg -> []
end
end,
Tables).
|