Add option to use prepared statements for SQLite PutMany

base/database/storage/sqlite/prepared.go

@@ -0,0 +1,125 @@
+package sqlite
+
+import (
+	"context"
+	"strconv"
+
+	"github.com/stephenafamo/bob"
+	"github.com/stephenafamo/bob/dialect/sqlite/im"
+	"github.com/stephenafamo/bob/expr"
+
+	"github.com/safing/portmaster/base/database/record"
+	"github.com/safing/portmaster/base/database/storage/sqlite/models"
+	"github.com/safing/structures/dsd"
+)
+
+var UsePreparedStatements bool = true
+
+// PutMany stores many records in the database.
+func (db *SQLite) putManyWithPreparedStmts(shadowDelete bool) (chan<- record.Record, <-chan error) {
+	batch := make(chan record.Record, 100)
+	errs := make(chan error, 1)
+
+	// Simulate upsert with custom selection on conflict.
+	rawQuery, _, err := models.Records.Insert(
+		im.Into("records", "key", "format", "value", "created", "modified", "expires", "deleted", "secret", "crownjewel"),
+		im.Values(expr.Arg("key"), expr.Arg("format"), expr.Arg("value"), expr.Arg("created"), expr.Arg("modified"), expr.Arg("expires"), expr.Arg("deleted"), expr.Arg("secret"), expr.Arg("crownjewel")),
+		im.OnConflict("key").DoUpdate(
+			im.SetExcluded("format", "value", "created", "modified", "expires", "deleted", "secret", "crownjewel"),
+		),
+	).Build(db.ctx)
+	if err != nil {
+		errs <- err
+		return batch, errs
+	}
+
+	// Start transaction.
+	tx, err := db.bob.BeginTx(db.ctx, nil)
+	if err != nil {
+		errs <- err
+		return batch, errs
+	}
+
+	// Create prepared statement WITHIN TRANSACTION.
+	preparedStmt, err := tx.PrepareContext(db.ctx, rawQuery)
+	if err != nil {
+		errs <- err
+		return batch, errs
+	}
+
+	// start handler
+	go func() {
+		// Read all put records.
+	writeBatch:
+		for {
+			select {
+			case r := <-batch:
+				if r != nil {
+					// Write record.
+					err := writeWithPreparedStatement(db.ctx, &preparedStmt, r)
+					if err != nil {
+						errs <- err
+						break writeBatch
+					}
+				} else {
+					// Finalize transcation.
+					errs <- tx.Commit()
+					return
+				}
+
+			case <-db.ctx.Done():
+				break writeBatch
+			}
+		}
+
+		// Rollback transaction.
+		errs <- tx.Rollback()
+	}()
+
+	return batch, errs
+}
+
+func writeWithPreparedStatement(ctx context.Context, pStmt *bob.StdPrepared, r record.Record) error {
+	r.Lock()
+	defer r.Unlock()
+
+	// Serialize to JSON.
+	data, err := r.MarshalDataOnly(r, dsd.JSON)
+	if err != nil {
+		return err
+	}
+
+	// Get Meta.
+	m := r.Meta()
+
+	// Insert.
+	if len(data) > 0 {
+		format := strconv.Itoa(dsd.JSON)
+		_, err = pStmt.ExecContext(
+			ctx,
+			r.DatabaseKey(),
+			format,
+			data,
+			m.Created,
+			m.Modified,
+			m.Expires,
+			m.Deleted,
+			m.IsSecret(),
+			m.IsCrownJewel(),
+		)
+	} else {
+		_, err = pStmt.ExecContext(
+			ctx,
+			r.DatabaseKey(),
+			nil,
+			nil,
+			m.Created,
+			m.Modified,
+			m.Expires,
+			m.Deleted,
+			m.IsSecret(),
+			m.IsCrownJewel(),
+		)
+	}
+	return err
+}

base/database/storage/sqlite/prepared_test.go

@@ -0,0 +1,86 @@
+package sqlite
+
+import (
+	"strconv"
+	"testing"
+)
+
+func BenchmarkPutMany(b *testing.B) {
+	// Configure prepared statement usage.
+	origSetting := UsePreparedStatements
+	UsePreparedStatements = false
+	defer func() {
+		UsePreparedStatements = origSetting
+	}()
+
+	// Run benchmark.
+	benchPutMany(b)
+}
+
+func BenchmarkPutManyPreparedStmt(b *testing.B) {
+	// Configure prepared statement usage.
+	origSetting := UsePreparedStatements
+	UsePreparedStatements = true
+	defer func() {
+		UsePreparedStatements = origSetting
+	}()
+
+	// Run benchmark.
+	benchPutMany(b)
+}
+
+func benchPutMany(b *testing.B) { //nolint:thelper
+	// Start database.
+	testDir := b.TempDir()
+	db, err := openSQLite("test", testDir, false)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer func() {
+		// shutdown
+		err = db.Shutdown()
+		if err != nil {
+			b.Fatal(err)
+		}
+	}()
+
+	// Start benchmarking.
+	b.ResetTimer()
+
+	// Benchmark PutMany.
+	records, errs := db.PutMany(false)
+	for i := range b.N {
+		// Create test record.
+		newTestRecord := &TestRecord{
+			S:    "banana",
+			I:    42,
+			I8:   42,
+			I16:  42,
+			I32:  42,
+			I64:  42,
+			UI:   42,
+			UI8:  42,
+			UI16: 42,
+			UI32: 42,
+			UI64: 42,
+			F32:  42.42,
+			F64:  42.42,
+			B:    true,
+		}
+		newTestRecord.UpdateMeta()
+		newTestRecord.SetKey("test:" + strconv.Itoa(i))
+
+		select {
+		case records <- newTestRecord:
+		case err := <-errs:
+			b.Fatal(err)
+		}
+	}
+
+	// Finalize.
+	close(records)
+	err = <-errs
+	if err != nil {
+		b.Fatal(err)
+	}
+}

base/database/storage/sqlite/sqlite.go

@@ -205,6 +205,11 @@ func (db *SQLite) PutMany(shadowDelete bool) (chan<- record.Record, <-chan error
 	db.wg.Add(1)
 	defer db.wg.Done()
 
+	// Check if we should use prepared statement optimized inserting.
+	if UsePreparedStatements {
+		return db.putManyWithPreparedStmts(shadowDelete)
+	}
+
 	batch := make(chan record.Record, 100)
 	errs := make(chan error, 1)