Move various limiter types behind a common interface

src/internal/common/limiters/token_bucket.go

@@ -0,0 +1,32 @@
+package limiters
+
+import (
+	"context"
+
+	"golang.org/x/time/rate"
+)
+
+var _ Limiter = &TokenBucket{}
+
+// Thin wrapper around the golang.org/x/time/rate token bucket rate limiter.
+type TokenBucket struct {
+	*rate.Limiter
+}
+
+func NewTokenBucketLimiter(r int, burst int) Limiter {
+	return &TokenBucket{
+		Limiter: rate.NewLimiter(rate.Limit(r), burst),
+	}
+}
+
+func (tb *TokenBucket) Wait(ctx context.Context) error {
+	return tb.Limiter.Wait(ctx)
+}
+
+func (tb *TokenBucket) WaitN(ctx context.Context, n int) error {
+	return tb.Limiter.WaitN(ctx, n)
+}
+
+// Reset and shutdown are no-ops for the token bucket limiter.
+func (tb *TokenBucket) Reset()    {}
+func (tb *TokenBucket) Shutdown() {}

src/pkg/services/m365/api/graph/concurrency_middleware.go

@@ -9,8 +9,8 @@ import (
 
 	"github.com/alcionai/clues"
 	khttp "github.com/microsoft/kiota-http-go"
-	"golang.org/x/time/rate"
 
+	"github.com/alcionai/corso/src/internal/common/limiters"
 	"github.com/alcionai/corso/src/pkg/count"
 	"github.com/alcionai/corso/src/pkg/logger"
 	"github.com/alcionai/corso/src/pkg/path"
@@ -99,9 +99,9 @@ const (
 )
 
 var (
-	driveLimiter = rate.NewLimiter(drivePerSecond, driveMaxCap)
+	driveLimiter = limiters.NewTokenBucketLimiter(drivePerSecond, driveMaxCap)
 	// also used as the exchange service limiter
-	defaultLimiter = rate.NewLimiter(defaultPerSecond, defaultMaxCap)
+	defaultLimiter = limiters.NewTokenBucketLimiter(defaultPerSecond, defaultMaxCap)
 )
 
 type LimiterCfg struct {
@@ -116,7 +116,7 @@ func BindRateLimiterConfig(ctx context.Context, lc LimiterCfg) context.Context {
 	return context.WithValue(ctx, limiterCfgCtxKey, lc)
 }
 
-func ctxLimiter(ctx context.Context) *rate.Limiter {
+func ctxLimiter(ctx context.Context) limiters.Limiter {
 	lc, ok := extractRateLimiterConfig(ctx)
 	if !ok {
 		return defaultLimiter

src/pkg/services/m365/api/graph/middleware_test.go

@@ -19,8 +19,8 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
-	"golang.org/x/time/rate"
 
+	"github.com/alcionai/corso/src/internal/common/limiters"
 	"github.com/alcionai/corso/src/internal/common/ptr"
 	"github.com/alcionai/corso/src/internal/tester"
 	"github.com/alcionai/corso/src/internal/tester/tconfig"
@@ -356,7 +356,7 @@ func (suite *MiddlewareUnitSuite) TestBindExtractLimiterConfig() {
 		name          string
 		service       path.ServiceType
 		expectOK      require.BoolAssertionFunc
-		expectLimiter *rate.Limiter
+		expectLimiter limiters.Limiter
 	}{
 		{
 			name:          "exchange",