From 05d5ec480497a0b0952eba2de5305e611cfb6d73 Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Fri, 29 Jan 2021 19:24:41 -0500
Subject: [PATCH] state: remove the need for registerSchema

registerSchema creates some indirection which is not necessary in this
case. newDBSchema can call each of the tables.

Enterprise tables can be added from the existing withEnterpriseSchema
shim.
---
 agent/consul/state/acl.go              |  8 ----
 agent/consul/state/autopilot.go        |  7 +--
 agent/consul/state/catalog_schema.go   |  8 ----
 agent/consul/state/config_entry.go     |  4 --
 agent/consul/state/connect_ca.go       |  6 ---
 agent/consul/state/coordinate.go       |  7 +--
 agent/consul/state/federation_state.go |  4 --
 agent/consul/state/intention.go        |  4 --
 agent/consul/state/kvs.go              |  8 +---
 agent/consul/state/prepared_query.go   |  7 +--
 agent/consul/state/schema.go           | 59 +++++++++++++++-----------
 agent/consul/state/schema_test.go      |  2 +-
 agent/consul/state/session.go          |  8 +---
 agent/consul/state/state_store.go      |  2 +-
 agent/consul/state/system_metadata.go  |  3 --
 agent/consul/state/usage.go            |  4 --
 16 files changed, 46 insertions(+), 95 deletions(-)

diff --git a/agent/consul/state/acl.go b/agent/consul/state/acl.go
index bff97c7189..7bd11194f6 100644
--- a/agent/consul/state/acl.go
+++ b/agent/consul/state/acl.go
@@ -212,14 +212,6 @@ func (s *TokenExpirationIndex) FromArgs(args ...interface{}) ([]byte, error) {
 	return buf, nil
 }
 
-func init() {
-	registerSchema(tokensTableSchema)
-	registerSchema(policiesTableSchema)
-	registerSchema(rolesTableSchema)
-	registerSchema(bindingRulesTableSchema)
-	registerSchema(authMethodsTableSchema)
-}
-
 // ACLTokens is used when saving a snapshot
 func (s *Snapshot) ACLTokens() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("acl-tokens", "id")
diff --git a/agent/consul/state/autopilot.go b/agent/consul/state/autopilot.go
index d1791f34a3..14c804b286 100644
--- a/agent/consul/state/autopilot.go
+++ b/agent/consul/state/autopilot.go
@@ -3,8 +3,9 @@ package state
 import (
 	"fmt"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // autopilotConfigTableSchema returns a new table schema used for storing
@@ -25,10 +26,6 @@ func autopilotConfigTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(autopilotConfigTableSchema)
-}
-
 // Autopilot is used to pull the autopilot config from the snapshot.
 func (s *Snapshot) Autopilot() (*structs.AutopilotConfig, error) {
 	c, err := s.tx.First("autopilot-config", "id")
diff --git a/agent/consul/state/catalog_schema.go b/agent/consul/state/catalog_schema.go
index 1c69cd6fa9..6b5ebfbf18 100644
--- a/agent/consul/state/catalog_schema.go
+++ b/agent/consul/state/catalog_schema.go
@@ -347,11 +347,3 @@ func (index *ServiceNameIndex) PrefixFromArgs(args ...interface{}) ([]byte, erro
 	}
 	return val, nil
 }
-
-func init() {
-	registerSchema(nodesTableSchema)
-	registerSchema(servicesTableSchema)
-	registerSchema(checksTableSchema)
-	registerSchema(gatewayServicesTableSchema)
-	registerSchema(meshTopologyTableSchema)
-}
diff --git a/agent/consul/state/config_entry.go b/agent/consul/state/config_entry.go
index 98ef58ed30..8fa28e9a2a 100644
--- a/agent/consul/state/config_entry.go
+++ b/agent/consul/state/config_entry.go
@@ -74,10 +74,6 @@ func (s *ConfigEntryLinkIndex) PrefixFromArgs(args ...interface{}) ([]byte, erro
 	return val, nil
 }
 
-func init() {
-	registerSchema(configTableSchema)
-}
-
 // ConfigEntries is used to pull all the config entries for the snapshot.
 func (s *Snapshot) ConfigEntries() ([]structs.ConfigEntry, error) {
 	entries, err := s.tx.Get(tableConfigEntries, "id")
diff --git a/agent/consul/state/connect_ca.go b/agent/consul/state/connect_ca.go
index aa5134fb1f..73c68f8afa 100644
--- a/agent/consul/state/connect_ca.go
+++ b/agent/consul/state/connect_ca.go
@@ -73,12 +73,6 @@ func caRootTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(caBuiltinProviderTableSchema)
-	registerSchema(caConfigTableSchema)
-	registerSchema(caRootTableSchema)
-}
-
 // CAConfig is used to pull the CA config from the snapshot.
 func (s *Snapshot) CAConfig() (*structs.CAConfiguration, error) {
 	c, err := s.tx.First(tableConnectCAConfig, "id")
diff --git a/agent/consul/state/coordinate.go b/agent/consul/state/coordinate.go
index 9b4bd3a96e..13271021db 100644
--- a/agent/consul/state/coordinate.go
+++ b/agent/consul/state/coordinate.go
@@ -3,9 +3,10 @@ package state
 import (
 	"fmt"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/consul/lib"
-	"github.com/hashicorp/go-memdb"
 )
 
 // coordinatesTableSchema returns a new table schema used for storing
@@ -47,10 +48,6 @@ func coordinatesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(coordinatesTableSchema)
-}
-
 // Coordinates is used to pull all the coordinates from the snapshot.
 func (s *Snapshot) Coordinates() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("coordinates", "id")
diff --git a/agent/consul/state/federation_state.go b/agent/consul/state/federation_state.go
index 5c93500730..058433f11d 100644
--- a/agent/consul/state/federation_state.go
+++ b/agent/consul/state/federation_state.go
@@ -27,10 +27,6 @@ func federationStateTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(federationStateTableSchema)
-}
-
 // FederationStates is used to pull all the federation states for the snapshot.
 func (s *Snapshot) FederationStates() ([]*structs.FederationState, error) {
 	configs, err := s.tx.Get(tableFederationStates, "id")
diff --git a/agent/consul/state/intention.go b/agent/consul/state/intention.go
index e367bcc09f..c227010bf2 100644
--- a/agent/consul/state/intention.go
+++ b/agent/consul/state/intention.go
@@ -95,10 +95,6 @@ func intentionsTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(intentionsTableSchema)
-}
-
 // LegacyIntentions is used to pull all the intentions from the snapshot.
 //
 // Deprecated: service-intentions config entries are handled as config entries
diff --git a/agent/consul/state/kvs.go b/agent/consul/state/kvs.go
index 554f3830be..db7d731e3e 100644
--- a/agent/consul/state/kvs.go
+++ b/agent/consul/state/kvs.go
@@ -4,8 +4,9 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // kvsTableSchema returns a new table schema used for storing key/value data for
@@ -48,11 +49,6 @@ func tombstonesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(kvsTableSchema)
-	registerSchema(tombstonesTableSchema)
-}
-
 // KVs is used to pull the full list of KVS entries for use during snapshots.
 func (s *Snapshot) KVs() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("kvs", "id_prefix")
diff --git a/agent/consul/state/prepared_query.go b/agent/consul/state/prepared_query.go
index 3009dc577d..28bead6901 100644
--- a/agent/consul/state/prepared_query.go
+++ b/agent/consul/state/prepared_query.go
@@ -4,9 +4,10 @@ import (
 	"fmt"
 	"regexp"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/hashicorp/consul/agent/consul/prepared_query"
 	"github.com/hashicorp/consul/agent/structs"
-	"github.com/hashicorp/go-memdb"
 )
 
 // preparedQueriesTableSchema returns a new table schema used for storing
@@ -50,10 +51,6 @@ func preparedQueriesTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(preparedQueriesTableSchema)
-}
-
 // validUUID is used to check if a given string looks like a UUID
 var validUUID = regexp.MustCompile(`(?i)^[\da-f]{8}-[\da-f]{4}-[\da-f]{4}-[\da-f]{4}-[\da-f]{12}$`)
 
diff --git a/agent/consul/state/schema.go b/agent/consul/state/schema.go
index 9cb0bb93d6..052e4aca88 100644
--- a/agent/consul/state/schema.go
+++ b/agent/consul/state/schema.go
@@ -6,28 +6,43 @@ import (
 	"github.com/hashicorp/go-memdb"
 )
 
-// schemaFn is an interface function used to create and return
-// new memdb schema structs for constructing an in-memory db.
-type schemaFn func() *memdb.TableSchema
+// newDBSchema creates and returns the memdb schema for the Store.
+func newDBSchema() *memdb.DBSchema {
+	db := &memdb.DBSchema{Tables: make(map[string]*memdb.TableSchema)}
 
-// schemas is used to register schemas with the state store.
-var schemas []schemaFn
-
-// registerSchema registers a new schema with the state store. This should
-// get called at package init() time.
-func registerSchema(fn schemaFn) {
-	schemas = append(schemas, fn)
+	addTableSchemas(db,
+		authMethodsTableSchema,
+		autopilotConfigTableSchema,
+		bindingRulesTableSchema,
+		caBuiltinProviderTableSchema,
+		caConfigTableSchema,
+		caRootTableSchema,
+		checksTableSchema,
+		configTableSchema,
+		coordinatesTableSchema,
+		federationStateTableSchema,
+		gatewayServicesTableSchema,
+		indexTableSchema,
+		intentionsTableSchema,
+		kvsTableSchema,
+		meshTopologyTableSchema,
+		nodesTableSchema,
+		policiesTableSchema,
+		preparedQueriesTableSchema,
+		rolesTableSchema,
+		servicesTableSchema,
+		sessionChecksTableSchema,
+		sessionsTableSchema,
+		systemMetadataTableSchema,
+		tokensTableSchema,
+		tombstonesTableSchema,
+		usageTableSchema,
+	)
+	withEnterpriseSchema(db)
+	return db
 }
 
-// stateStoreSchema is used to return the combined schema for
-// the state store.
-func stateStoreSchema() *memdb.DBSchema {
-	// Create the root DB schema
-	db := &memdb.DBSchema{
-		Tables: make(map[string]*memdb.TableSchema),
-	}
-
-	// Add the tables to the root schema
+func addTableSchemas(db *memdb.DBSchema, schemas ...func() *memdb.TableSchema) {
 	for _, fn := range schemas {
 		schema := fn()
 		if _, ok := db.Tables[schema.Name]; ok {
@@ -35,8 +50,6 @@ func stateStoreSchema() *memdb.DBSchema {
 		}
 		db.Tables[schema.Name] = schema
 	}
-	withEnterpriseSchema(db)
-	return db
 }
 
 // indexTableSchema returns a new table schema used for tracking various indexes
@@ -57,7 +70,3 @@ func indexTableSchema() *memdb.TableSchema {
 		},
 	}
 }
-
-func init() {
-	registerSchema(indexTableSchema)
-}
diff --git a/agent/consul/state/schema_test.go b/agent/consul/state/schema_test.go
index 0d9a47838f..8fff28c05e 100644
--- a/agent/consul/state/schema_test.go
+++ b/agent/consul/state/schema_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestStateStoreSchema(t *testing.T) {
-	schema := stateStoreSchema()
+	schema := newDBSchema()
 	require.NoError(t, schema.Validate())
 
 	_, err := memdb.NewMemDB(schema)
diff --git a/agent/consul/state/session.go b/agent/consul/state/session.go
index 36dce655eb..dd6815ae67 100644
--- a/agent/consul/state/session.go
+++ b/agent/consul/state/session.go
@@ -6,8 +6,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/go-memdb"
+
+	"github.com/hashicorp/consul/agent/structs"
 )
 
 // sessionsTableSchema returns a new table schema used for storing session
@@ -129,11 +130,6 @@ func (index *CheckIDIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
 	return val, nil
 }
 
-func init() {
-	registerSchema(sessionsTableSchema)
-	registerSchema(sessionChecksTableSchema)
-}
-
 // Sessions is used to pull the full list of sessions for use during snapshots.
 func (s *Snapshot) Sessions() (memdb.ResultIterator, error) {
 	iter, err := s.tx.Get("sessions", "id")
diff --git a/agent/consul/state/state_store.go b/agent/consul/state/state_store.go
index 26a7c76a2a..c38ddfa8d8 100644
--- a/agent/consul/state/state_store.go
+++ b/agent/consul/state/state_store.go
@@ -155,7 +155,7 @@ type sessionCheck struct {
 // NewStateStore creates a new in-memory state storage layer.
 func NewStateStore(gc *TombstoneGC) *Store {
 	// Create the in-memory DB.
-	schema := stateStoreSchema()
+	schema := newDBSchema()
 	db, err := memdb.NewMemDB(schema)
 	if err != nil {
 		// the only way for NewMemDB to error is if the schema is invalid. The
diff --git a/agent/consul/state/system_metadata.go b/agent/consul/state/system_metadata.go
index 0655758a60..09ed0202dc 100644
--- a/agent/consul/state/system_metadata.go
+++ b/agent/consul/state/system_metadata.go
@@ -26,9 +26,6 @@ func systemMetadataTableSchema() *memdb.TableSchema {
 		},
 	}
 }
-func init() {
-	registerSchema(systemMetadataTableSchema)
-}
 
 // SystemMetadataEntries used to pull all the system metadata entries for the snapshot.
 func (s *Snapshot) SystemMetadataEntries() ([]*structs.SystemMetadataEntry, error) {
diff --git a/agent/consul/state/usage.go b/agent/consul/state/usage.go
index 459de23fc4..f9720b3e0f 100644
--- a/agent/consul/state/usage.go
+++ b/agent/consul/state/usage.go
@@ -31,10 +31,6 @@ func usageTableSchema() *memdb.TableSchema {
 	}
 }
 
-func init() {
-	registerSchema(usageTableSchema)
-}
-
 // UsageEntry represents a count of some arbitrary identifier within the
 // state store, along with the last seen index.
 type UsageEntry struct {