From 6bc80ca3bea7b79eea1d4621a876760c74f5768d Mon Sep 17 00:00:00 2001
From: Abhishek Pandey <abhishek@alcion.ai>
Date: Wed, 15 Nov 2023 19:49:48 -0800
Subject: [PATCH] Move rate limiter ctx binding to backup operation scope
 (#4679)

Currently we specify the service to rate limiter type bindings in `ProduceBackupCollections`. This means we lose the binding during `ConsumeBackupCollections` and default to using the exchange token bucket limiter for all services. Onedrive & SP are an exception here since we bind again during `streamItems`.

This PR moves the binding call one level above, to `do` scope so that the same bindings stay applied for the entirety of the backup op.

Note: This change has a slight side effect on integration tests which directly test `ProduceBackupCollections`. We'll default to using exchange token bucket limiter for such integ tests for all services. This shouldn't impact test runs too much since exch limts(16 per sec) or upto 1160 per min are still pretty good compared to 1200 per min for onedrive. Also, since these tests are short running & won't be doing a load test, we should be fine here.

If there is an objection, I can add rate limiter bindings manually for these `ProduceBackupCollections` tests.

---

#### Does this PR need a docs update or release note?

- [ ] :white_check_mark: Yes, it's included
- [ ] :clock1: Yes, but in a later PR
- [x] :no_entry: No

#### Type of change

<!--- Please check the type of change your PR introduces: --->
- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [ ] :robot: Supportability/Tests
- [ ] :computer: CI/Deployment
- [x] :broom: Tech Debt/Cleanup

#### Issue(s)

<!-- Can reference multiple issues. Use one of the following "magic words" - "closes, fixes" to auto-close the Github issue. -->
* #<issue>

#### Test Plan

<!-- How will this be tested prior to merging.-->
- [x] :muscle: Manual
- [ ] :zap: Unit test
- [ ] :green_heart: E2E
---
 src/internal/m365/backup.go                      | 2 --
 src/internal/m365/collection/drive/collection.go | 8 --------
 src/internal/operations/backup.go                | 7 +++++++
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/src/internal/m365/backup.go b/src/internal/m365/backup.go
index a09b11b0e..10e29dab1 100644
--- a/src/internal/m365/backup.go
+++ b/src/internal/m365/backup.go
@@ -46,8 +46,6 @@ func (ctrl *Controller) ProduceBackupCollections(
 		diagnostics.Index("service", bpc.Selector.PathService().String()))
 	defer end()
 
-	ctx = graph.BindRateLimiterConfig(ctx, graph.LimiterCfg{Service: service})
-
 	// Limit the max number of active requests to graph from this collection.
 	bpc.Options.Parallelism.ItemFetch = graph.Parallelism(service).
 		ItemOverride(ctx, bpc.Options.Parallelism.ItemFetch)
diff --git a/src/internal/m365/collection/drive/collection.go b/src/internal/m365/collection/drive/collection.go
index 00f99abe2..d06859450 100644
--- a/src/internal/m365/collection/drive/collection.go
+++ b/src/internal/m365/collection/drive/collection.go
@@ -587,14 +587,6 @@ func (oc *Collection) streamDriveItem(
 		parentPath)
 
 	ctx = clues.Add(ctx, "item_info", itemInfo)
-	// Drive content download requests are also rate limited by graph api.
-	// Ensure that this request goes through the drive limiter & not the default
-	// limiter.
-	ctx = graph.BindRateLimiterConfig(
-		ctx,
-		graph.LimiterCfg{
-			Service: path.OneDriveService,
-		})
 
 	if isFile {
 		dataSuffix := metadata.DataFileSuffix
diff --git a/src/internal/operations/backup.go b/src/internal/operations/backup.go
index e6e5a5490..4b9699e7a 100644
--- a/src/internal/operations/backup.go
+++ b/src/internal/operations/backup.go
@@ -205,6 +205,13 @@ func (op *BackupOperation) Run(ctx context.Context) (err error) {
 
 	ctx = clues.AddTrace(ctx)
 
+	// Select an appropriate rate limiter for the service.
+	ctx = graph.BindRateLimiterConfig(
+		ctx,
+		graph.LimiterCfg{
+			Service: op.Selectors.PathService(),
+		})
+
 	// Check if the protected resource has the service enabled in order for us
 	// to run a backup.
 	enabled, err := op.bp.IsServiceEnabled(