From 6aae21541f2e534acee2a8568f4b1b0ecfc057b1 Mon Sep 17 00:00:00 2001
From: JT Olio <hello@jtolio.com>
Date: Sat, 28 Nov 2020 13:54:52 -0700
Subject: [PATCH] satellitedb: do saverollup in batches

Change-Id: I78278a192cba60541eee2986f54a88d5a479bd3e
---
 cmd/satellite/main.go                         |  1 +
 satellite/accounting/tally/tally.go           |  3 +-
 satellite/satellitedb/database.go             |  3 ++
 .../satellitedb/storagenodeaccounting.go      | 49 ++++++++++++++-----
 scripts/testdata/satellite-config.yaml.lock   |  3 ++
 5 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/cmd/satellite/main.go b/cmd/satellite/main.go
index a09fb759f..a669bff89 100644
--- a/cmd/satellite/main.go
+++ b/cmd/satellite/main.go
@@ -327,6 +327,7 @@ func cmdRun(cmd *cobra.Command, args []string) (err error) {
 
 	db, err := satellitedb.Open(ctx, log.Named("db"), runCfg.Database, satellitedb.Options{
 		ReportedRollupsReadBatchSize: runCfg.Orders.SettlementBatchSize,
+		SaveRollupBatchSize:          runCfg.Tally.SaveRollupBatchSize,
 	})
 	if err != nil {
 		return errs.New("Error starting master database on satellite: %+v", err)
diff --git a/satellite/accounting/tally/tally.go b/satellite/accounting/tally/tally.go
index cebd06e22..2cb179562 100644
--- a/satellite/accounting/tally/tally.go
+++ b/satellite/accounting/tally/tally.go
@@ -27,7 +27,8 @@ var (
 
 // Config contains configurable values for the tally service.
 type Config struct {
-	Interval time.Duration `help:"how frequently the tally service should run" releaseDefault:"1h" devDefault:"30s"`
+	Interval            time.Duration `help:"how frequently the tally service should run" releaseDefault:"1h" devDefault:"30s"`
+	SaveRollupBatchSize int           `help:"how large of batches SaveRollup should process at a time" default:"1000"`
 }
 
 // Service is the tally service for data stored on each storage node
diff --git a/satellite/satellitedb/database.go b/satellite/satellitedb/database.go
index 48efda99a..29e0984b1 100644
--- a/satellite/satellitedb/database.go
+++ b/satellite/satellitedb/database.go
@@ -73,6 +73,9 @@ type Options struct {
 	// How many records to read in a single transaction when asked for all of the
 	// billable bandwidth from the reported serials table.
 	ReportedRollupsReadBatchSize int
+
+	// How many storage node rollups to save in one batch.
+	SaveRollupBatchSize int
 }
 
 var _ dbx.DBMethods = &satelliteDB{}
diff --git a/satellite/satellitedb/storagenodeaccounting.go b/satellite/satellitedb/storagenodeaccounting.go
index 3db38b709..159f21383 100644
--- a/satellite/satellitedb/storagenodeaccounting.go
+++ b/satellite/satellitedb/storagenodeaccounting.go
@@ -152,9 +152,26 @@ func (db *StoragenodeAccounting) SaveRollup(ctx context.Context, latestRollup ti
 	if len(stats) == 0 {
 		return Error.New("In SaveRollup with empty nodeData")
 	}
-	err = db.db.WithTx(ctx, func(ctx context.Context, tx *dbx.Tx) error {
-		for _, arsByDate := range stats {
-			for _, ar := range arsByDate {
+
+	batchSize := db.db.opts.SaveRollupBatchSize
+	if batchSize <= 0 {
+		batchSize = 1000
+	}
+
+	var rollups []*accounting.Rollup
+	for _, arsByDate := range stats {
+		for _, ar := range arsByDate {
+			rollups = append(rollups, ar)
+		}
+	}
+	finished := false
+
+	for !finished {
+		err = db.db.WithTx(ctx, func(ctx context.Context, tx *dbx.Tx) error {
+			for i := 0; i < batchSize && len(rollups) > 0; i++ {
+				ar := rollups[0]
+				rollups = rollups[1:]
+
 				nID := dbx.AccountingRollup_NodeId(ar.NodeID.Bytes())
 				start := dbx.AccountingRollup_StartTime(ar.StartTime)
 				put := dbx.AccountingRollup_PutTotal(ar.PutTotal)
@@ -169,16 +186,24 @@ func (db *StoragenodeAccounting) SaveRollup(ctx context.Context, latestRollup ti
 					return err
 				}
 			}
-		}
 
-		return tx.UpdateNoReturn_AccountingTimestamps_By_Name(ctx,
-			dbx.AccountingTimestamps_Name(accounting.LastRollup),
-			dbx.AccountingTimestamps_Update_Fields{
-				Value: dbx.AccountingTimestamps_Value(latestRollup),
-			},
-		)
-	})
-	return Error.Wrap(err)
+			if len(rollups) == 0 {
+				finished = true
+				return tx.UpdateNoReturn_AccountingTimestamps_By_Name(ctx,
+					dbx.AccountingTimestamps_Name(accounting.LastRollup),
+					dbx.AccountingTimestamps_Update_Fields{
+						Value: dbx.AccountingTimestamps_Value(latestRollup),
+					},
+				)
+			}
+
+			return nil
+		})
+		if err != nil {
+			return Error.Wrap(err)
+		}
+	}
+	return nil
 }
 
 // LastTimestamp records the greatest last tallied time.
diff --git a/scripts/testdata/satellite-config.yaml.lock b/scripts/testdata/satellite-config.yaml.lock
index 8eda15ab0..001a2a7a1 100755
--- a/scripts/testdata/satellite-config.yaml.lock
+++ b/scripts/testdata/satellite-config.yaml.lock
@@ -667,6 +667,9 @@ server.private-address: 127.0.0.1:7778
 # how frequently the tally service should run
 # tally.interval: 1h0m0s
 
+# how large of batches SaveRollup should process at a time
+# tally.save-rollup-batch-size: 1000
+
 # address for jaeger agent
 # tracing.agent-addr: agent.tracing.datasci.storj.io:5775