Skip to content

Commit

Permalink
refactor(client): take need Limiter
Browse files Browse the repository at this point in the history
  • Loading branch information
syd committed Jul 29, 2022
1 parent 9551df6 commit 0f03434
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 11 deletions.
9 changes: 2 additions & 7 deletions client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,7 @@ func (c *Client) Close() {

func (c *Client) NewLimiter(ctx context.Context, name string, rate int, interval time.Duration) error {
clt := rpc.GetClientFromPool(c.pool)
req := new(pb.Limiter)
req.Name = name
req.Rate = int32(rate)
req.Interval = int64(interval)
req := pb.NewLimiter(name, rate, interval)
return clt.Call(ctx, "NewLimiter", req, nil)
}

Expand All @@ -45,10 +42,8 @@ func (c *Client) RemoveLimiter(ctx context.Context, name string) error {
return clt.Call(ctx, "RemoveLimiter", req, nil)
}

func (c *Client) Take(ctx context.Context, name string) (time.Duration, error) {
func (c *Client) Take(ctx context.Context, req *pb.Limiter) (time.Duration, error) {
clt := rpc.GetClientFromPool(c.pool)
req := new(pb.Limiter)
req.Name = name
ret := new(pb.Limiter)
err := clt.Call(ctx, "Take", req, ret)
return time.Duration(ret.Interval), err
Expand Down
11 changes: 11 additions & 0 deletions pb/limiter.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
package pb

import "time"

func NewLimiter(name string, rate int, interval time.Duration) *Limiter {
ret := new(Limiter)
ret.Name = name
ret.Rate = int32(rate)
ret.Interval = interval.Nanoseconds()
return ret
}
9 changes: 5 additions & 4 deletions service/rpc/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,11 +48,12 @@ func (s *Service) Take(ctx context.Context, req *pb.Limiter, ret *pb.Limiter) er
s.lock.RLock()
limiter, found := s.mp[req.GetName()]
s.lock.RUnlock()
if found {
now := time.Now()
t := limiter.Take()
ret.Interval = t.Sub(now).Nanoseconds()
if !found && req.GetRate() > 0 && req.GetInterval() > 0 {
s.NewLimiter(ctx, req, nil)
}
now := time.Now()
t := limiter.Take()
ret.Interval = t.Sub(now).Nanoseconds()
return nil
}

Expand Down

0 comments on commit 0f03434

Please sign in to comment.