Abstracting out Table implementation and supporting Swiss and Go maps

internal/common/map.go

@@ -0,0 +1,9 @@
+package common
+
+type ITable[K comparable, V any] interface {
+	Put(key K, value V)
+	Get(key K) (V, bool)
+	Delete(key K)
+	Len() int
+	All(func(k K, obj V) bool)
+}

internal/common/regmap.go

@@ -0,0 +1,30 @@
+package common
+
+type RegMap[K comparable, V any] struct {
+	M map[K]V
+}
+
+func (t *RegMap[K, V]) Put(key K, value V) {
+	t.M[key] = value
+}
+
+func (t *RegMap[K, V]) Get(key K) (V, bool) {
+	value, ok := t.M[key]
+	return value, ok
+}
+
+func (t *RegMap[K, V]) Delete(key K) {
+	delete(t.M, key)
+}
+
+func (t *RegMap[K, V]) Len() int {
+	return len(t.M)
+}
+
+func (t *RegMap[K, V]) All(f func(k K, obj V) bool) {
+	for k, v := range t.M {
+		if !f(k, v) {
+			break
+		}
+	}
+}

internal/common/swisstable.go

@@ -0,0 +1,27 @@
+package common
+
+import "github.com/cockroachdb/swiss"
+
+type SwissTable[K comparable, V any] struct {
+	M *swiss.Map[K, V]
+}
+
+func (t *SwissTable[K, V]) Put(key K, value V) {
+	t.M.Put(key, value)
+}
+
+func (t *SwissTable[K, V]) Get(key K) (V, bool) {
+	return t.M.Get(key)
+}
+
+func (t *SwissTable[K, V]) Delete(key K) {
+	t.M.Delete(key)
+}
+
+func (t *SwissTable[K, V]) Len() int {
+	return t.M.Len()
+}
+
+func (t *SwissTable[K, V]) All(f func(k K, obj V) bool) {
+	t.M.All(f)
+}

internal/querymanager/query_manager.go

@@ -10,20 +10,20 @@
 	"time"
 
 	"github.com/dicedb/dice/internal/comm"
+	"github.com/dicedb/dice/internal/common"
 
 	"github.com/ohler55/ojg/jp"
 
 	"github.com/dicedb/dice/internal/object"
 
 	"github.com/dicedb/dice/internal/sql"
 
-	"github.com/cockroachdb/swiss"
 	"github.com/dicedb/dice/internal/clientio"
 	dstore "github.com/dicedb/dice/internal/store"
 )
 
 type (
-	cacheStore *swiss.Map[string, *object.Obj]
+	cacheStore common.ITable[string, *object.Obj]
 
 	// QuerySubscription represents a subscription to watch a query.
 	QuerySubscription struct {
@@ -53,8 +53,8 @@
 
 	// Manager watches for changes in keys and notifies clients.
 	Manager struct {
-		WatchList    sync.Map                       // WatchList is a map of query string to their respective clients, type: map[string]*sync.Map[int]struct{}
-		QueryCache   *swiss.Map[string, cacheStore] // QueryCache is a map of fingerprints to their respective data caches
+		WatchList    sync.Map                          // WatchList is a map of query string to their respective clients, type: map[string]*sync.Map[int]struct{}
+		QueryCache   common.ITable[string, cacheStore] // QueryCache is a map of fingerprints to their respective data caches
 		QueryCacheMu sync.RWMutex
 		logger       *slog.Logger
 	}
@@ -86,21 +86,37 @@
 	}
 }
 
+func NewQueryCacheStoreRegMap() common.ITable[string, cacheStore] {
+	return &common.RegMap[string, cacheStore]{
+		M: make(map[string]cacheStore),
+	}
+}
+
+func NewQueryCacheStore() common.ITable[string, cacheStore] {
+	return NewQueryCacheStoreRegMap()
+}
+
+func NewCacheStoreRegMap() cacheStore {
+	return &common.RegMap[string, *object.Obj]{
+		M: make(map[string]*object.Obj),
+	}
+}
+
+func NewCacheStore() cacheStore {
+	return NewCacheStoreRegMap()
+}
+
 // NewQueryManager initializes a new Manager.
 func NewQueryManager(logger *slog.Logger) *Manager {
 	QuerySubscriptionChan = make(chan QuerySubscription)
 	AdhocQueryChan = make(chan AdhocQuery, 1000)
 	return &Manager{
 		WatchList:  sync.Map{},
-		QueryCache: swiss.New[string, cacheStore](0),
+		QueryCache: NewQueryCacheStore(),
 		logger:     logger,
 	}
 }
 
-func newCacheStore() cacheStore {
-	return swiss.New[string, *object.Obj](0)
-}
-
 // Run starts the Manager's main loops.
 func (m *Manager) Run(ctx context.Context, watchChan <-chan dstore.QueryWatchEvent) {
 	var wg sync.WaitGroup
@@ -211,9 +227,9 @@
 
 	switch event.Operation {
 	case dstore.Set:
-		((*swiss.Map[string, *object.Obj])(store)).Put(event.Key, &event.Value)
+		store.Put(event.Key, &event.Value)
 	case dstore.Del:
-		((*swiss.Map[string, *object.Obj])(store)).Delete(event.Key)
+		store.Delete(event.Key)
 	default:
 		m.logger.Warn("Unknown operation", slog.String("operation", event.Operation))
 	}
@@ -314,13 +330,13 @@
 	m.QueryCacheMu.Lock()
 	defer m.QueryCacheMu.Unlock()
 
-	cache := newCacheStore()
+	cache := NewCacheStore()
 	// Hydrate the cache with data from all shards.
 	// TODO: We need to ensure we receive cache data from all shards once we have multithreading in place.
 	//  For now we only expect one update.
 	kvs := <-cacheChan
 	for _, kv := range *kvs {
-		((*swiss.Map[string, *object.Obj])(cache)).Put(kv.Key, kv.Value)
+		cache.Put(kv.Key, kv.Value)
 	}
 
 	m.QueryCache.Put(query.Fingerprint, cache)

internal/sql/executor.go

@@ -10,7 +10,7 @@ import (
 	"github.com/dicedb/dice/internal/object"
 
 	"github.com/bytedance/sonic"
-	"github.com/cockroachdb/swiss"
+	"github.com/dicedb/dice/internal/common"
 	"github.com/dicedb/dice/internal/regex"
 	"github.com/dicedb/dice/internal/server/utils"
 	"github.com/ohler55/ojg/jp"
@@ -31,7 +31,7 @@ type QueryResultRowWithOrder struct {
 	OrderByType  string
 }
 
-func ExecuteQuery(query *DSQLQuery, store *swiss.Map[string, *object.Obj]) ([]QueryResultRow, error) {
+func ExecuteQuery(query *DSQLQuery, store common.ITable[string, *object.Obj]) ([]QueryResultRow, error) {
 	var result []QueryResultRow
 	var err error
 	jsonPathCache := make(map[string]jp.Expr) // Cache for parsed JSON paths

internal/store/expire_test.go

@@ -1,17 +1,16 @@
 package store
 
 import (
-	"github.com/dicedb/dice/internal/object"
 	"testing"
 
-	"github.com/cockroachdb/swiss"
+	"github.com/dicedb/dice/internal/object"
 )
 
 func TestDelExpiry(t *testing.T) {
 	store := NewStore(nil)
 	// Initialize the test environment
-	store.store = swiss.New[string, *object.Obj](10240)
-	store.expires = swiss.New[*object.Obj, uint64](10240)
+	store.store = NewStoreMap()
+	store.expires = NewExpireMap()
 
 	// Define test cases
 	tests := []struct {

internal/store/store.go

@@ -5,16 +5,36 @@ import (
 
 	"github.com/ohler55/ojg/jp"
 
+	"github.com/dicedb/dice/internal/common"
 	"github.com/dicedb/dice/internal/object"
 	"github.com/dicedb/dice/internal/sql"
 	"github.com/xwb1989/sqlparser"
 
 	"github.com/dicedb/dice/internal/server/utils"
 
-	"github.com/cockroachdb/swiss"
 	"github.com/dicedb/dice/config"
 )
 
+func NewStoreRegMap() common.ITable[string, *object.Obj] {
+	return &common.RegMap[string, *object.Obj]{
+		M: make(map[string]*object.Obj),
+	}
+}
+
+func NewExpireRegMap() common.ITable[*object.Obj, uint64] {
+	return &common.RegMap[*object.Obj, uint64]{
+		M: make(map[*object.Obj]uint64),
+	}
+}
+
+func NewStoreMap() common.ITable[string, *object.Obj] {
+	return NewStoreRegMap()
+}
+
+func NewExpireMap() common.ITable[*object.Obj, uint64] {
+	return NewExpireRegMap()
+}
+
 // QueryWatchEvent represents a change in a watched key.
 type QueryWatchEvent struct {
 	Key       string
@@ -23,24 +43,24 @@ type QueryWatchEvent struct {
 }
 
 type Store struct {
-	store     *swiss.Map[string, *object.Obj]
-	expires   *swiss.Map[*object.Obj, uint64] // Does not need to be thread-safe as it is only accessed by a single thread.
+	store     common.ITable[string, *object.Obj]
+	expires   common.ITable[*object.Obj, uint64] // Does not need to be thread-safe as it is only accessed by a single thread.
 	numKeys   int
 	watchChan chan QueryWatchEvent
 }
 
 func NewStore(watchChan chan QueryWatchEvent) *Store {
 	return &Store{
-		store:     swiss.New[string, *object.Obj](config.DiceConfig.Server.StoreMapInitSize),
-		expires:   swiss.New[*object.Obj, uint64](config.DiceConfig.Server.StoreMapInitSize),
+		store:     NewStoreRegMap(),
+		expires:   NewExpireRegMap(),
 		watchChan: watchChan,
 	}
 }
 
 func ResetStore(store *Store) *Store {
 	store.numKeys = 0
-	store.store = swiss.New[string, *object.Obj](config.DiceConfig.Server.StoreMapInitSize)
-	store.expires = swiss.New[*object.Obj, uint64](config.DiceConfig.Server.StoreMapInitSize)
+	store.store = NewStoreMap()
+	store.expires = NewExpireMap()
 
 	return store
 }
@@ -59,8 +79,8 @@ func (store *Store) NewObj(value interface{}, expDurationMs int64, oType, oEnc u
 
 func (store *Store) ResetStore() {
 	store.numKeys = 0
-	store.store.Clear()
-	store.expires.Clear()
+	store.store = NewStoreMap()
+	store.expires = NewExpireMap()
 }
 
 type PutOptions struct {
@@ -295,7 +315,7 @@ func (store *Store) notifyQueryManager(k, operation string, obj object.Obj) {
 	store.watchChan <- QueryWatchEvent{k, operation, obj}
 }
 
-func (store *Store) GetStore() *swiss.Map[string, *object.Obj] {
+func (store *Store) GetStore() common.ITable[string, *object.Obj] {
 	return store.store
 }