From 9aad0aa42876e22ece793bd1e3322406c4537f90 Mon Sep 17 00:00:00 2001
From: Abhishek Pandey <abhishek@alcion.ai>
Date: Tue, 26 Sep 2023 16:40:48 +0530
Subject: [PATCH] Add initial stats

---
 src/cli/backup/backup.go                      |  2 ++
 .../m365/graph/concurrency_middleware.go      |  2 ++
 src/internal/m365/graph/middleware.go         |  6 ++++++
 src/internal/operations/backup.go             |  5 +++++
 src/internal/stats/stats.go                   | 20 +++++++++++++++++++
 src/pkg/count/keys.go                         |  3 ++-
 6 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/cli/backup/backup.go b/src/cli/backup/backup.go
index 2d3db4597..52ad3741d 100644
--- a/src/cli/backup/backup.go
+++ b/src/cli/backup/backup.go
@@ -214,6 +214,8 @@ func genericCreateCommand(
 			continue
 		}
 
+		logger.Ctx(ictx).Infow("graph api stats", "stats", bo.Results.APIStats)
+
 		bIDs = append(bIDs, string(bo.Results.BackupID))
 
 		if !DisplayJSONFormat() {
diff --git a/src/internal/m365/graph/concurrency_middleware.go b/src/internal/m365/graph/concurrency_middleware.go
index c470fc1cf..d93bd2246 100644
--- a/src/internal/m365/graph/concurrency_middleware.go
+++ b/src/internal/m365/graph/concurrency_middleware.go
@@ -191,6 +191,8 @@ func QueueRequest(ctx context.Context) {
 	if err := limiter.WaitN(ctx, consume); err != nil {
 		logger.CtxErr(ctx, err).Error("graph middleware waiting on the limiter")
 	}
+
+	// Add to the count bus
 }
 
 // RateLimiterMiddleware is used to ensure we don't overstep per-min request limits.
diff --git a/src/internal/m365/graph/middleware.go b/src/internal/m365/graph/middleware.go
index d5935a5b0..4adf8cef8 100644
--- a/src/internal/m365/graph/middleware.go
+++ b/src/internal/m365/graph/middleware.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/alcionai/clues"
@@ -325,6 +326,8 @@ func (mw RetryMiddleware) getRetryDelay(
 // MetricsMiddleware aggregates per-request metrics on the events bus
 type MetricsMiddleware struct{}
 
+var xmCount int64
+
 const xmruHeader = "x-ms-resource-unit"
 
 func (mw *MetricsMiddleware) Intercept(
@@ -362,6 +365,9 @@ func (mw *MetricsMiddleware) Intercept(
 		xmrui = 1
 	}
 
+	atomic.AddInt64(&xmCount, int64(xmrui))
+	logger.Ctx(req.Context()).Info("xmcount ", xmCount)
+
 	countBus := count.Ctx(req.Context())
 	countBus.Add(count.APICallTokensConsumed, int64(xmrui))
 
diff --git a/src/internal/operations/backup.go b/src/internal/operations/backup.go
index 5635bc6f4..e9bfc7021 100644
--- a/src/internal/operations/backup.go
+++ b/src/internal/operations/backup.go
@@ -70,6 +70,7 @@ type BackupResults struct {
 	BackupID model.StableID `json:"backupID"`
 	// keys are found in /pkg/count/keys.go
 	Counts map[string]int64 `json:"counts"`
+	stats.APIStats
 }
 
 // NewBackupOperation constructs and validates a backup operation.
@@ -846,6 +847,10 @@ func (op *BackupOperation) persistResults(
 
 	op.Results.ItemsRead = opStats.ctrl.Successes
 
+	// API stats
+	apiStats := stats.GetAPIStats(op.Counter)
+	op.Results.TokensConsumed = apiStats.TokensConsumed
+
 	// Only return non-recoverable errors at this point.
 	return op.Errors.Failure()
 }
diff --git a/src/internal/stats/stats.go b/src/internal/stats/stats.go
index 4ecc5a811..bd790dc9f 100644
--- a/src/internal/stats/stats.go
+++ b/src/internal/stats/stats.go
@@ -3,6 +3,8 @@ package stats
 import (
 	"sync/atomic"
 	"time"
+
+	"github.com/alcionai/corso/src/pkg/count"
 )
 
 // ReadWrites tracks the total count of reads and writes.  ItemsRead
@@ -43,3 +45,21 @@ type SkippedCounts struct {
 	SkippedMalware            int `json:"skippedMalware"`
 	SkippedInvalidOneNoteFile int `json:"skippedInvalidOneNoteFile"`
 }
+
+type APIStats struct {
+	TokensConsumed int64 `json:"tokensConsumed"`
+	// PeakTokenUsage is the maximum number of tokens used during a
+	// rolling 1 minute window.
+	PeakTokenUsagePerMin int64 `json:"peakTokenUsage"`
+}
+
+func GetAPIStats(
+	ctr *count.Bus,
+) APIStats {
+	s := APIStats{}
+
+	s.TokensConsumed = ctr.Total(count.APICallTokensConsumed)
+	s.PeakTokenUsagePerMin = ctr.Get(count.PeakAPITokenUsagePerMin)
+
+	return s
+}
diff --git a/src/pkg/count/keys.go b/src/pkg/count/keys.go
index bdfd03ca7..90cc9e7b0 100644
--- a/src/pkg/count/keys.go
+++ b/src/pkg/count/keys.go
@@ -4,7 +4,8 @@ type key string
 
 const (
 	// count of bucket-tokens consumed by api calls.
-	APICallTokensConsumed key = "api-call-tokens-consumed"
+	APICallTokensConsumed   key = "api-call-tokens-consumed"
+	PeakAPITokenUsagePerMin key = "peak-api-token-usage-per-min"
 	// count of api calls that resulted in failure due to throttling.
 	ThrottledAPICalls key = "throttled-api-calls"
 )