ratelimit.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. /*
  2. Copyright 2018 Google LLC
  3. Licensed under the Apache License, Version 2.0 (the "License");
  4. you may not use this file except in compliance with the License.
  5. You may obtain a copy of the License at
  6. https://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. package cloud
  14. import (
  15. "context"
  16. "time"
  17. "github.com/GoogleCloudPlatform/k8s-cloud-provider/pkg/cloud/meta"
  18. )
  19. // RateLimitKey is a key identifying the operation to be rate limited. The rate limit
  20. // queue will be determined based on the contents of RateKey.
  21. type RateLimitKey struct {
  22. // ProjectID is the non-numeric ID of the project.
  23. ProjectID string
  24. // Operation is the specific method being invoked (e.g. "Get", "List").
  25. Operation string
  26. // Version is the API version of the call.
  27. Version meta.Version
  28. // Service is the service being invoked (e.g. "Firewalls", "BackendServices")
  29. Service string
  30. }
  31. // RateLimiter is the interface for a rate limiting policy.
  32. type RateLimiter interface {
  33. // Accept uses the RateLimitKey to derive a sleep time for the calling
  34. // goroutine. This call will block until the operation is ready for
  35. // execution.
  36. //
  37. // Accept returns an error if the given context ctx was canceled
  38. // while waiting for acceptance into the queue.
  39. Accept(ctx context.Context, key *RateLimitKey) error
  40. }
  41. // acceptor is an object which blocks within Accept until a call is allowed to run.
  42. // Accept is a behavior of the flowcontrol.RateLimiter interface.
  43. type acceptor interface {
  44. // Accept blocks until a call is allowed to run.
  45. Accept()
  46. }
  47. // AcceptRateLimiter wraps an Acceptor with RateLimiter parameters.
  48. type AcceptRateLimiter struct {
  49. // Acceptor is the underlying rate limiter.
  50. Acceptor acceptor
  51. }
  52. // Accept wraps an Acceptor and blocks on Accept or context.Done(). Key is ignored.
  53. func (rl *AcceptRateLimiter) Accept(ctx context.Context, key *RateLimitKey) error {
  54. ch := make(chan struct{})
  55. go func() {
  56. rl.Acceptor.Accept()
  57. close(ch)
  58. }()
  59. select {
  60. case <-ch:
  61. break
  62. case <-ctx.Done():
  63. return ctx.Err()
  64. }
  65. return nil
  66. }
  67. // NopRateLimiter is a rate limiter that performs no rate limiting.
  68. type NopRateLimiter struct {
  69. }
  70. // Accept everything immediately.
  71. func (*NopRateLimiter) Accept(ctx context.Context, key *RateLimitKey) error {
  72. return nil
  73. }
  74. // MinimumRateLimiter wraps a RateLimiter and will only call its Accept until the minimum
  75. // duration has been met or the context is cancelled.
  76. type MinimumRateLimiter struct {
  77. // RateLimiter is the underlying ratelimiter which is called after the mininum time is reacehd.
  78. RateLimiter RateLimiter
  79. // Minimum is the minimum wait time before the underlying ratelimiter is called.
  80. Minimum time.Duration
  81. }
  82. // Accept blocks on the minimum duration and context. Once the minimum duration is met,
  83. // the func is blocked on the underlying ratelimiter.
  84. func (m *MinimumRateLimiter) Accept(ctx context.Context, key *RateLimitKey) error {
  85. select {
  86. case <-time.After(m.Minimum):
  87. return m.RateLimiter.Accept(ctx, key)
  88. case <-ctx.Done():
  89. return ctx.Err()
  90. }
  91. }