@@ -36,18 +36,14 @@ func VirtualNetwork(rs mockrouting.Server, d delay.D) Network {
36
36
}
37
37
}
38
38
39
- type rateLimiter interface {
40
- Limit (dataSize int ) time.Duration
41
- }
42
-
43
39
type RateLimitGenerator interface {
44
40
NextRateLimit () float64
45
41
}
46
42
47
43
func RateLimitedVirtualNetwork (rs mockrouting.Server , d delay.D , rateLimitGenerator RateLimitGenerator ) Network {
48
44
return & network {
49
45
latencies : make (map [peer.ID ]map [peer.ID ]time.Duration ),
50
- rateLimiters : make (map [peer.ID ]map [peer.ID ]rateLimiter ),
46
+ rateLimiters : make (map [peer.ID ]map [peer.ID ]* mocknet. RateLimiter ),
51
47
clients : make (map [peer.ID ]* receiverQueue ),
52
48
delay : d ,
53
49
routingserver : rs ,
@@ -60,7 +56,7 @@ func RateLimitedVirtualNetwork(rs mockrouting.Server, d delay.D, rateLimitGenera
60
56
type network struct {
61
57
mu sync.Mutex
62
58
latencies map [peer.ID ]map [peer.ID ]time.Duration
63
- rateLimiters map [peer.ID ]map [peer.ID ]rateLimiter
59
+ rateLimiters map [peer.ID ]map [peer.ID ]* mocknet. RateLimiter
64
60
clients map [peer.ID ]* receiverQueue
65
61
routingserver mockrouting.Server
66
62
delay delay.D
@@ -133,18 +129,18 @@ func (n *network) SendMessage(
133
129
if n .isRateLimited {
134
130
rateLimiters , ok := n .rateLimiters [from ]
135
131
if ! ok {
136
- rateLimiters = make (map [peer.ID ]rateLimiter )
132
+ rateLimiters = make (map [peer.ID ]* mocknet. RateLimiter )
137
133
n .rateLimiters [from ] = rateLimiters
138
134
}
139
135
140
- rl , ok := rateLimiters [to ]
136
+ rateLimiter , ok := rateLimiters [to ]
141
137
if ! ok {
142
- rl = mocknet .NewRatelimiter (n .rateLimitGenerator .NextRateLimit ())
143
- rateLimiters [to ] = rl
138
+ rateLimiter = mocknet .NewRateLimiter (n .rateLimitGenerator .NextRateLimit ())
139
+ rateLimiters [to ] = rateLimiter
144
140
}
145
141
146
142
size := mes .ToProtoV1 ().Size ()
147
- bandwidthDelay = rl .Limit (size )
143
+ bandwidthDelay = rateLimiter .Limit (size )
148
144
} else {
149
145
bandwidthDelay = 0
150
146
}
0 commit comments