From cd9dfafd2d3bd523a70d4cdc0325982b9d30f8f7 Mon Sep 17 00:00:00 2001
From: Ivan Valdes <ivan@vald.es>
Date: Sun, 11 Feb 2024 20:25:57 -0800
Subject: [PATCH] server: Implement WithMmapSize for backend config

Revert to receiving the log and the path. Accept a third argument which
are options to be applied to the backend config.
Add a new funcion WithMmapSize, which modifies the backend config to
override this option.

Signed-off-by: Ivan Valdes <ivan@vald.es>
---
 etcdutl/etcdutl/migrate_command.go            |  2 +-
 etcdutl/snapshot/v3_snapshot.go               |  8 ++++----
 server/etcdserver/cindex/cindex_test.go       |  4 ++--
 server/storage/backend/backend.go             | 20 +++++++++++--------
 server/storage/mvcc/store_test.go             |  4 ++--
 server/storage/schema/auth_roles_test.go      |  4 ++--
 server/storage/schema/auth_test.go            |  4 ++--
 server/storage/schema/auth_users_test.go      |  4 ++--
 server/storage/schema/lease_test.go           |  2 +-
 server/storage/schema/schema_test.go          |  6 +++---
 server/storage/schema/version_test.go         |  2 +-
 server/verify/verify.go                       |  3 ++-
 tests/e2e/utl_migrate_test.go                 |  2 +-
 .../integration/clientv3/maintenance_test.go  |  2 +-
 tests/integration/v3_alarm_test.go            |  2 +-
 tools/etcd-dump-db/backend.go                 |  2 +-
 16 files changed, 38 insertions(+), 33 deletions(-)

diff --git a/etcdutl/etcdutl/migrate_command.go b/etcdutl/etcdutl/migrate_command.go
index 45ea407e48b7..521cf8ba80c8 100644
--- a/etcdutl/etcdutl/migrate_command.go
+++ b/etcdutl/etcdutl/migrate_command.go
@@ -92,7 +92,7 @@ func (o *migrateOptions) Config() (*migrateConfig, error) {
 	}
 
 	dbPath := datadir.ToBackendFileName(o.dataDir)
-	c.be = backend.NewDefaultBackend(backend.BackendConfig{Logger: GetLogger(), Path: dbPath})
+	c.be = backend.NewDefaultBackend(GetLogger(), dbPath)
 
 	walPath := datadir.ToWalDir(o.dataDir)
 	w, err := wal.OpenForRead(c.lg, walPath, walpb.Snapshot{})
diff --git a/etcdutl/snapshot/v3_snapshot.go b/etcdutl/snapshot/v3_snapshot.go
index 4f02ce09719e..b892047ae265 100644
--- a/etcdutl/snapshot/v3_snapshot.go
+++ b/etcdutl/snapshot/v3_snapshot.go
@@ -325,7 +325,7 @@ func (s *v3Manager) saveDB() error {
 		return err
 	}
 
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: s.lg, Path: s.outDbPath(), MmapSize: s.initialMmapSize})
+	be := backend.NewDefaultBackend(s.lg, s.outDbPath(), backend.WithMmapSize(s.initialMmapSize))
 	defer be.Close()
 
 	err = schema.NewMembershipBackend(s.lg, be).TrimMembershipFromBackend()
@@ -339,7 +339,7 @@ func (s *v3Manager) saveDB() error {
 // modifyLatestRevision can increase the latest revision by the given amount and sets the scheduled compaction
 // to that revision so that the server will consider this revision compacted.
 func (s *v3Manager) modifyLatestRevision(bumpAmount uint64) error {
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: s.lg, Path: s.outDbPath()})
+	be := backend.NewDefaultBackend(s.lg, s.outDbPath())
 	defer func() {
 		be.ForceCommit()
 		be.Close()
@@ -480,7 +480,7 @@ func (s *v3Manager) saveWALAndSnap() (*raftpb.HardState, error) {
 	}
 
 	// add members again to persist them to the store we create.
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: s.lg, Path: s.outDbPath(), MmapSize: s.initialMmapSize})
+	be := backend.NewDefaultBackend(s.lg, s.outDbPath(), backend.WithMmapSize(s.initialMmapSize))
 	defer be.Close()
 	s.cl.SetBackend(schema.NewMembershipBackend(s.lg, be))
 	for _, m := range s.cl.Members() {
@@ -559,7 +559,7 @@ func (s *v3Manager) saveWALAndSnap() (*raftpb.HardState, error) {
 }
 
 func (s *v3Manager) updateCIndex(commit uint64, term uint64) error {
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: s.lg, Path: s.outDbPath(), MmapSize: s.initialMmapSize})
+	be := backend.NewDefaultBackend(s.lg, s.outDbPath(), backend.WithMmapSize(s.initialMmapSize))
 	defer be.Close()
 
 	cindex.UpdateConsistentIndexForce(be.BatchTx(), commit, term)
diff --git a/server/etcdserver/cindex/cindex_test.go b/server/etcdserver/cindex/cindex_test.go
index 2f78daf0f1d5..a056ac3d759c 100644
--- a/server/etcdserver/cindex/cindex_test.go
+++ b/server/etcdserver/cindex/cindex_test.go
@@ -56,7 +56,7 @@ func TestConsistentIndex(t *testing.T) {
 	be.ForceCommit()
 	be.Close()
 
-	b := backend.NewDefaultBackend(backend.BackendConfig{Logger: zaptest.NewLogger(t), Path: tmpPath})
+	b := backend.NewDefaultBackend(zaptest.NewLogger(t), tmpPath)
 	defer b.Close()
 	ci.SetBackend(b)
 	index = ci.ConsistentIndex()
@@ -108,7 +108,7 @@ func TestConsistentIndexDecrease(t *testing.T) {
 			be.ForceCommit()
 			be.Close()
 
-			be = backend.NewDefaultBackend(backend.BackendConfig{Logger: zaptest.NewLogger(t), Path: tmpPath})
+			be = backend.NewDefaultBackend(zaptest.NewLogger(t), tmpPath)
 			defer be.Close()
 			ci := NewConsistentIndex(be)
 			ci.SetConsistentIndex(tc.index, tc.term)
diff --git a/server/storage/backend/backend.go b/server/storage/backend/backend.go
index 217809bd256c..aa5e70faba9a 100644
--- a/server/storage/backend/backend.go
+++ b/server/storage/backend/backend.go
@@ -151,6 +151,8 @@ type BackendConfig struct {
 	Hooks Hooks
 }
 
+type BackendConfigFunc func(*BackendConfig)
+
 func DefaultBackendConfig(lg *zap.Logger) BackendConfig {
 	return BackendConfig{
 		BatchInterval: defaultBatchInterval,
@@ -164,15 +166,17 @@ func New(bcfg BackendConfig) Backend {
 	return newBackend(bcfg)
 }
 
-func NewDefaultBackend(bc BackendConfig) Backend {
-
-	bcfg := DefaultBackendConfig(bc.Logger)
-	bcfg.Path = bc.Path
+func WithMmapSize(size uint64) BackendConfigFunc {
+	return func(bcfg *BackendConfig) {
+		bcfg.MmapSize = size
+	}
+}
 
-	if bc.MmapSize > 0 {
-		bcfg.MmapSize = bc.MmapSize
-	} else {
-		bcfg.MmapSize = InitialMmapSize
+func NewDefaultBackend(lg *zap.Logger, path string, bcfns ...BackendConfigFunc) Backend {
+	bcfg := DefaultBackendConfig(lg)
+	bcfg.Path = path
+	for _, fn := range bcfns {
+		fn(&bcfg)
 	}
 
 	return newBackend(bcfg)
diff --git a/server/storage/mvcc/store_test.go b/server/storage/mvcc/store_test.go
index 3abcb0c76b8a..bd6d25e171fa 100644
--- a/server/storage/mvcc/store_test.go
+++ b/server/storage/mvcc/store_test.go
@@ -61,7 +61,7 @@ func TestScheduledCompact(t *testing.T) {
 			be.ForceCommit()
 			be.Close()
 
-			b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			b := backend.NewDefaultBackend(lg, tmpPath)
 			defer b.Close()
 			v, found := UnsafeReadScheduledCompact(b.BatchTx())
 			assert.Equal(t, true, found)
@@ -103,7 +103,7 @@ func TestFinishedCompact(t *testing.T) {
 			be.ForceCommit()
 			be.Close()
 
-			b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			b := backend.NewDefaultBackend(lg, tmpPath)
 			defer b.Close()
 			v, found := UnsafeReadFinishedCompact(b.BatchTx())
 			assert.Equal(t, true, found)
diff --git a/server/storage/schema/auth_roles_test.go b/server/storage/schema/auth_roles_test.go
index d86babc0708c..d7323c6f4a6e 100644
--- a/server/storage/schema/auth_roles_test.go
+++ b/server/storage/schema/auth_roles_test.go
@@ -122,7 +122,7 @@ func TestGetAllRoles(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			users := abe2.GetAllRoles()
@@ -219,7 +219,7 @@ func TestGetRole(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			users := abe2.GetRole("role1")
diff --git a/server/storage/schema/auth_test.go b/server/storage/schema/auth_test.go
index 0f1277297165..96174e50ffc6 100644
--- a/server/storage/schema/auth_test.go
+++ b/server/storage/schema/auth_test.go
@@ -66,7 +66,7 @@ func TestAuthEnabled(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			tx = abe2.BatchTx()
@@ -117,7 +117,7 @@ func TestAuthRevision(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			tx := abe2.BatchTx()
diff --git a/server/storage/schema/auth_users_test.go b/server/storage/schema/auth_users_test.go
index 5911b795075d..2261e57071b0 100644
--- a/server/storage/schema/auth_users_test.go
+++ b/server/storage/schema/auth_users_test.go
@@ -110,7 +110,7 @@ func TestGetAllUsers(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			users := abe2.ReadTx().UnsafeGetAllUsers()
@@ -195,7 +195,7 @@ func TestGetUser(t *testing.T) {
 			abe.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			abe2 := NewAuthBackend(lg, be2)
 			users := abe2.GetUser("alice")
diff --git a/server/storage/schema/lease_test.go b/server/storage/schema/lease_test.go
index 68ffab204670..88a1cd7e0633 100644
--- a/server/storage/schema/lease_test.go
+++ b/server/storage/schema/lease_test.go
@@ -99,7 +99,7 @@ func TestLeaseBackend(t *testing.T) {
 			be.ForceCommit()
 			be.Close()
 
-			be2 := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			be2 := backend.NewDefaultBackend(lg, tmpPath)
 			defer be2.Close()
 			leases := MustUnsafeGetAllLeases(be2.ReadTx())
 
diff --git a/server/storage/schema/schema_test.go b/server/storage/schema/schema_test.go
index 312286c4959d..3d9bd1065c3a 100644
--- a/server/storage/schema/schema_test.go
+++ b/server/storage/schema/schema_test.go
@@ -83,7 +83,7 @@ func TestValidate(t *testing.T) {
 			lg := zap.NewNop()
 			dataPath := setupBackendData(t, tc.version, tc.overrideKeys)
 
-			b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: dataPath})
+			b := backend.NewDefaultBackend(lg, dataPath)
 			defer b.Close()
 			err := Validate(lg, b.ReadTx())
 			if (err != nil) != tc.expectError {
@@ -211,7 +211,7 @@ func TestMigrate(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: dataPath})
+			b := backend.NewDefaultBackend(lg, dataPath)
 			defer b.Close()
 
 			err = Migrate(lg, b.BatchTx(), walVersion, tc.targetVersion)
@@ -255,7 +255,7 @@ func TestMigrateIsReversible(t *testing.T) {
 			lg := zap.NewNop()
 			dataPath := setupBackendData(t, tc.initialVersion, nil)
 
-			be := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: dataPath})
+			be := backend.NewDefaultBackend(lg, dataPath)
 			defer be.Close()
 			tx := be.BatchTx()
 			tx.Lock()
diff --git a/server/storage/schema/version_test.go b/server/storage/schema/version_test.go
index 59ccd774dc4a..576f063df777 100644
--- a/server/storage/schema/version_test.go
+++ b/server/storage/schema/version_test.go
@@ -69,7 +69,7 @@ func TestVersion(t *testing.T) {
 			be.ForceCommit()
 			be.Close()
 
-			b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: tmpPath})
+			b := backend.NewDefaultBackend(lg, tmpPath)
 			defer b.Close()
 			v := UnsafeReadStorageVersion(b.BatchTx())
 
diff --git a/server/verify/verify.go b/server/verify/verify.go
index 036a7c647396..6aaf1422aee5 100644
--- a/server/verify/verify.go
+++ b/server/verify/verify.go
@@ -74,7 +74,8 @@ func Verify(cfg Config) error {
 			lg.Info("verification of persisted state successful", zap.String("data-dir", cfg.DataDir))
 		}
 	}()
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: datadir.ToBackendFileName(cfg.DataDir)})
+
+	be := backend.NewDefaultBackend(lg, datadir.ToBackendFileName(cfg.DataDir))
 	defer be.Close()
 
 	snapshot, hardstate, err := validateWal(cfg)
diff --git a/tests/e2e/utl_migrate_test.go b/tests/e2e/utl_migrate_test.go
index 35d0b95fa94d..a58d69a24c74 100644
--- a/tests/e2e/utl_migrate_test.go
+++ b/tests/e2e/utl_migrate_test.go
@@ -166,7 +166,7 @@ func TestEtctlutlMigrate(t *testing.T) {
 			}
 
 			t.Log("etcdutl migrate...")
-			be := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: filepath.Join(memberDataDir, "member/snap/db")})
+			be := backend.NewDefaultBackend(lg, filepath.Join(memberDataDir, "member/snap/db"))
 			defer be.Close()
 
 			ver := schema.ReadStorageVersion(be.ReadTx())
diff --git a/tests/integration/clientv3/maintenance_test.go b/tests/integration/clientv3/maintenance_test.go
index c99f7846eb4d..8aa275bb90c4 100644
--- a/tests/integration/clientv3/maintenance_test.go
+++ b/tests/integration/clientv3/maintenance_test.go
@@ -277,7 +277,7 @@ func testMaintenanceSnapshotErrorInflight(t *testing.T, snapshot func(context.Co
 	// take about 1-second to read snapshot
 	clus.Members[0].Stop(t)
 	dpath := filepath.Join(clus.Members[0].DataDir, "member", "snap", "db")
-	b := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: dpath})
+	b := backend.NewDefaultBackend(lg, dpath)
 	s := mvcc.NewStore(lg, b, &lease.FakeLessor{}, mvcc.StoreConfig{CompactionBatchLimit: math.MaxInt32})
 	rev := 100000
 	for i := 2; i <= rev; i++ {
diff --git a/tests/integration/v3_alarm_test.go b/tests/integration/v3_alarm_test.go
index 8d9cd2b9d78b..50a701e68095 100644
--- a/tests/integration/v3_alarm_test.go
+++ b/tests/integration/v3_alarm_test.go
@@ -194,7 +194,7 @@ func TestV3CorruptAlarm(t *testing.T) {
 	// Corrupt member 0 by modifying backend offline.
 	clus.Members[0].Stop(t)
 	fp := filepath.Join(clus.Members[0].DataDir, "member", "snap", "db")
-	be := backend.NewDefaultBackend(backend.BackendConfig{Logger: lg, Path: fp})
+	be := backend.NewDefaultBackend(lg, fp)
 	s := mvcc.NewStore(lg, be, nil, mvcc.StoreConfig{})
 	// NOTE: cluster_proxy mode with namespacing won't set 'k', but namespace/'k'.
 	s.Put([]byte("abc"), []byte("def"), 0)
diff --git a/tools/etcd-dump-db/backend.go b/tools/etcd-dump-db/backend.go
index 37aa5b9125eb..5c9df5e42e05 100644
--- a/tools/etcd-dump-db/backend.go
+++ b/tools/etcd-dump-db/backend.go
@@ -179,7 +179,7 @@ func iterateBucket(dbPath, bucket string, limit uint64, decode bool) (err error)
 }
 
 func getHash(dbPath string) (hash uint32, err error) {
-	b := backend.NewDefaultBackend(backend.BackendConfig{Logger: zap.NewNop(), Path: dbPath})
+	b := backend.NewDefaultBackend(zap.NewNop(), dbPath)
 	return b.Hash(schema.DefaultIgnores)
 }