fix: batch sqlite store IO to better handle large deletes in migrations

Author: garethgeorge

internal/ioutil/iobatching.go

@@ -0,0 +1,15 @@
+package ioutil
+
+const DefaultBatchSize = 512
+
+func Batchify[T any](items []T, batchSize int) [][]T {
+	var batches [][]T
+	for i := 0; i < len(items); i += batchSize {
+		end := i + batchSize
+		if end > len(items) {
+			end = len(items)
+		}
+		batches = append(batches, items[i:end])
+	}
+	return batches
+}

internal/oplog/migrations.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	v1 "github.com/garethgeorge/backrest/gen/go/v1"
+	"github.com/garethgeorge/backrest/internal/ioutil"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/proto"
 )
@@ -120,8 +121,13 @@ func migration003DeduplicateIndexedSnapshots(oplog *OpLog) error {
 	if len(deleteIDs) == 0 {
 		return nil
 	}
-	_, err := oplog.store.Delete(deleteIDs...)
-	return err
+
+	for _, batch := range ioutil.Batchify(deleteIDs, ioutil.DefaultBatchSize) {
+		if _, err := oplog.store.Delete(batch...); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 // migrationNoop is a migration that does nothing; replaces deprecated migrations.

internal/oplog/sqlitestore/sqlitestore.go

@@ -13,6 +13,7 @@ import (
 
 	v1 "github.com/garethgeorge/backrest/gen/go/v1"
 	"github.com/garethgeorge/backrest/internal/cryptoutil"
+	"github.com/garethgeorge/backrest/internal/ioutil"
 	"github.com/garethgeorge/backrest/internal/oplog"
 	"github.com/garethgeorge/backrest/internal/protoutil"
 	lru "github.com/hashicorp/golang-lru/v2"
@@ -490,50 +491,68 @@ func (m *SqliteStore) Delete(opID ...int64) ([]*v1.Operation, error) {
 
 	ops := make([]*v1.Operation, 0, len(opID))
 	return ops, withImmediateSqliteTransaction(conn, func() error {
-		// fetch all the operations we're about to delete
-		predicate := []string{"operations.id IN ("}
-		args := []any{}
-		for i, id := range opID {
-			if i > 0 {
-				predicate = append(predicate, ",")
+		for _, batch := range ioutil.Batchify(opID, ioutil.DefaultBatchSize) {
+			// Optimize for the case of 1 element or batch size elements (which will be common)
+			useTransient := len(batch) != ioutil.DefaultBatchSize || len(batch) == 1
+			batchOps, err := m.deleteHelper(conn, useTransient, batch...)
+			if err != nil {
+				return err
 			}
-			predicate = append(predicate, "?")
-			args = append(args, id)
+			ops = append(ops, batchOps...)
 		}
-		predicate = append(predicate, ")")
-		predicateStr := strings.Join(predicate, "")
-
-		if err := sqlitex.ExecuteTransient(conn, "SELECT operations.operation FROM operations JOIN operation_groups ON operations.ogid = operation_groups.ogid WHERE "+predicateStr, &sqlitex.ExecOptions{
-			Args: args,
-			ResultFunc: func(stmt *sqlite.Stmt) error {
-				opBytes := make([]byte, stmt.ColumnLen(0))
-				n := stmt.GetBytes("operation", opBytes)
-				opBytes = opBytes[:n]
+		return nil
+	})
+}
 
-				var op v1.Operation
-				if err := proto.Unmarshal(opBytes, &op); err != nil {
-					return fmt.Errorf("unmarshal operation bytes: %v", err)
-				}
-				ops = append(ops, &op)
-				return nil
-			},
-		}); err != nil {
-			return fmt.Errorf("load operations for delete: %v", err)
+func (m *SqliteStore) deleteHelper(conn *sqlite.Conn, transient bool, opID ...int64) ([]*v1.Operation, error) {
+	// fetch all the operations we're about to delete
+	predicate := []string{"operations.id IN ("}
+	args := []any{}
+	for i, id := range opID {
+		if i > 0 {
+			predicate = append(predicate, ",")
 		}
+		predicate = append(predicate, "?")
+		args = append(args, id)
+	}
+	predicate = append(predicate, ")")
+	predicateStr := strings.Join(predicate, "")
 
-		if len(ops) != len(opID) {
-			return fmt.Errorf("couldn't find all operations to delete: %w", oplog.ErrNotExist)
-		}
+	var ops []*v1.Operation
+	if err := sqlitex.ExecuteTransient(conn, "SELECT operations.operation FROM operations JOIN operation_groups ON operations.ogid = operation_groups.ogid WHERE "+predicateStr, &sqlitex.ExecOptions{
+		Args: args,
+		ResultFunc: func(stmt *sqlite.Stmt) error {
+			opBytes := make([]byte, stmt.ColumnLen(0))
+			n := stmt.GetBytes("operation", opBytes)
+			opBytes = opBytes[:n]
 
-		// Delete the operations
-		if err := sqlitex.ExecuteTransient(conn, "DELETE FROM operations WHERE "+predicateStr, &sqlitex.ExecOptions{
-			Args: args,
-		}); err != nil {
-			return fmt.Errorf("delete operations: %v", err)
-		}
+			var op v1.Operation
+			if err := proto.Unmarshal(opBytes, &op); err != nil {
+				return fmt.Errorf("unmarshal operation bytes: %v", err)
+			}
+			ops = append(ops, &op)
+			return nil
+		},
+	}); err != nil {
+		return nil, fmt.Errorf("load operations for delete: %v", err)
+	}
 
-		return nil
-	})
+	if len(ops) != len(opID) {
+		return nil, fmt.Errorf("couldn't find all operations to delete: %w", oplog.ErrNotExist)
+	}
+
+	// Delete the operations
+	execFunc := sqlitex.Execute
+	if transient {
+		execFunc = sqlitex.ExecuteTransient
+	}
+	if err := execFunc(conn, "DELETE FROM operations WHERE "+predicateStr, &sqlitex.ExecOptions{
+		Args: args,
+	}); err != nil {
+		return nil, fmt.Errorf("delete operations: %v", err)
+	}
+
+	return ops, nil
 }
 
 func (m *SqliteStore) ResetForTest(t *testing.T) error {

internal/oplog/storetests/storecontract_test.go

@@ -651,6 +651,102 @@ func TestTransform(t *testing.T) {
 	}
 }
 
+func TestDelete(t *testing.T) {
+	t.Parallel()
+	for name, store := range StoresForTest(t) {
+		t.Run(name, func(t *testing.T) {
+			log, err := oplog.NewOpLog(store)
+			if err != nil {
+				t.Fatalf("error creating oplog: %v", err)
+			}
+
+			op := &v1.Operation{
+				UnixTimeStartMs: 1234,
+				PlanId:          "plan1",
+				RepoId:          "repo1",
+				RepoGuid:        "repo1",
+				InstanceId:      "instance1",
+				Op:              &v1.Operation_OperationBackup{},
+			}
+
+			if err := log.Add(op); err != nil {
+				t.Fatalf("error adding operation: %s", err)
+			}
+
+			if err := log.Delete(op.Id); err != nil {
+				t.Fatalf("error deleting operation: %s", err)
+			}
+
+			var ops []*v1.Operation
+			if err := log.Query(oplog.Query{}, func(op *v1.Operation) error {
+				ops = append(ops, op)
+				return nil
+			}); err != nil {
+				t.Fatalf("error querying operations: %s", err)
+			}
+
+			if len(ops) != 0 {
+				t.Errorf("expected 0 operations after deletion, got %d", len(ops))
+			}
+		})
+	}
+}
+
+func TestBulkDelete(t *testing.T) {
+	t.Parallel()
+	for name, store := range StoresForTest(t) {
+		t.Run(name, func(t *testing.T) {
+			log, err := oplog.NewOpLog(store)
+			if err != nil {
+				t.Fatalf("error creating oplog: %v", err)
+			}
+
+			// Add 2000 operations
+			var ops []*v1.Operation
+			for i := 0; i < 2000; i++ {
+				op := &v1.Operation{
+					UnixTimeStartMs: 1234,
+					PlanId:          fmt.Sprintf("plan%d", i),
+					RepoId:          fmt.Sprintf("repo%d", i),
+					RepoGuid:        fmt.Sprintf("repo%d", i),
+					InstanceId:      fmt.Sprintf("instance%d", i),
+					Op:              &v1.Operation_OperationBackup{},
+				}
+				ops = append(ops, op)
+			}
+
+			var ids []int64
+			if err := log.Add(ops...); err != nil {
+				t.Fatalf("error adding operations: %s", err)
+			}
+			for _, op := range ops {
+				ids = append(ids, op.Id)
+			}
+
+			// Delete all operations
+			err = log.Delete(ids...)
+			if err != nil {
+				t.Fatalf("error deleting operations: %s", err)
+			}
+			if len(ids) != 2000 {
+				t.Errorf("expected 2000 deleted operations, got %d", len(ids))
+			}
+
+			// Verify deletion
+			var count int
+			if err := log.Query(oplog.Query{}, func(op *v1.Operation) error {
+				count++
+				return nil
+			}); err != nil {
+				t.Fatalf("error querying operations: %s", err)
+			}
+			if count != 0 {
+				t.Errorf("expected 0 operations after deletion, got %d", count)
+			}
+		})
+	}
+}
+
 func TestQueryMetadata(t *testing.T) {
 	t.Parallel()
 	for name, store := range StoresForTest(t) {