perf: improve locking code

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - fix: doc comments from object fields should be present in generated GraphQL schema [#630](https://github.com/hypermodeinc/modus/pull/630)
 - feat: add neo4j support in modus [#636](https://github.com/hypermodeinc/modus/pull/636)
+- perf: improve locking code [#637](https://github.com/hypermodeinc/modus/pull/637)
 
 ## UNRELEASED - Go SDK
 

cspell.json

@@ -94,6 +94,7 @@
     "ldflags",
     "legacymodels",
     "Lessable",
+    "lestrrat",
     "linkname",
     "logit",
     "logits",
@@ -143,6 +144,7 @@
     "prereleases",
     "promhttp",
     "ptrs",
+    "puzpuzpuz",
     "quickstart",
     "reindex",
     "renameio",
@@ -203,6 +205,7 @@
     "Weasley",
     "Weaviate",
     "wundergraph",
+    "xsync",
     "xxhash",
     "zenquotes",
     "zerolog"

runtime/collections/collection.go

@@ -11,66 +11,50 @@ package collections
 
 import (
 	"fmt"
-	"sync"
 
 	"github.com/hypermodeinc/modus/runtime/collections/index/interfaces"
+	"github.com/puzpuzpuz/xsync/v3"
 )
 
 type collection struct {
-	collectionNamespaceMap map[string]interfaces.CollectionNamespace
-	mu                     sync.RWMutex
+	collectionNamespaceMap *xsync.MapOf[string, interfaces.CollectionNamespace]
 }
 
 func newCollection() *collection {
 	return &collection{
-		collectionNamespaceMap: map[string]interfaces.CollectionNamespace{},
+		collectionNamespaceMap: xsync.NewMapOf[string, interfaces.CollectionNamespace](),
 	}
 }
 
 func (c *collection) getCollectionNamespaceMap() map[string]interfaces.CollectionNamespace {
-	return c.collectionNamespaceMap
+	m := make(map[string]interfaces.CollectionNamespace, c.collectionNamespaceMap.Size())
+	c.collectionNamespaceMap.Range(func(key string, value interfaces.CollectionNamespace) bool {
+		m[key] = value
+		return true
+	})
+
+	return m
 }
 
 func (c *collection) findNamespace(namespace string) (interfaces.CollectionNamespace, error) {
-	c.mu.RLock()
-	defer c.mu.RUnlock()
-
-	ns, found := c.collectionNamespaceMap[namespace]
+	ns, found := c.collectionNamespaceMap.Load(namespace)
 	if !found {
 		return nil, errNamespaceNotFound
 	}
 	return ns, nil
 }
 
 func (c *collection) findOrCreateNamespace(namespace string, index interfaces.CollectionNamespace) (interfaces.CollectionNamespace, error) {
-	c.mu.RLock()
-	ns, found := c.collectionNamespaceMap[namespace]
-	if found {
-		defer c.mu.RUnlock()
-		return ns, nil
-	}
-
-	c.mu.RUnlock()
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	ns, found = c.collectionNamespaceMap[namespace]
-	if found {
-		return ns, nil
-	}
-
-	c.collectionNamespaceMap[namespace] = index
-	return index, nil
+	result, _ := c.collectionNamespaceMap.LoadOrStore(namespace, index)
+	return result, nil // TODO: remove unused error
 }
 
 func (c *collection) createCollectionNamespace(namespace string, index interfaces.CollectionNamespace) (interfaces.CollectionNamespace, error) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	if _, found := c.collectionNamespaceMap[namespace]; found {
+	_, found := c.collectionNamespaceMap.Load(namespace)
+	if found {
 		return nil, fmt.Errorf("namespace with name %s already exists", namespace)
 	}
 
-	c.collectionNamespaceMap[namespace] = index
+	c.collectionNamespaceMap.Store(namespace, index)
 	return index, nil
 }

runtime/collections/factory.go

@@ -19,6 +19,7 @@ import (
 	"github.com/hypermodeinc/modus/runtime/collections/index/interfaces"
 	"github.com/hypermodeinc/modus/runtime/db"
 	"github.com/hypermodeinc/modus/runtime/logger"
+	"github.com/puzpuzpuz/xsync/v3"
 )
 
 const collectionFactoryWriteInterval = 1
@@ -40,7 +41,7 @@ func newCollectionFactory() *collectionFactory {
 	return &collectionFactory{
 		collectionMap: map[string]*collection{
 			"": {
-				collectionNamespaceMap: map[string]interfaces.CollectionNamespace{},
+				collectionNamespaceMap: xsync.NewMapOf[string, interfaces.CollectionNamespace](),
 			},
 		},
 		quit: make(chan struct{}),
@@ -86,7 +87,7 @@ func (cf *collectionFactory) readFromPostgres(ctx context.Context) bool {
 	resetTimerFaster := false
 	var err error
 	for _, namespaceCollectionFactory := range cf.collectionMap {
-		for _, col := range namespaceCollectionFactory.collectionNamespaceMap {
+		for _, col := range namespaceCollectionFactory.getCollectionNamespaceMap() {
 			resetTimerFaster, err = loadTextsIntoCollection(ctx, col)
 			if err != nil {
 				logger.Err(ctx, err).

runtime/collections/in_mem/hnsw/vector_index.go

@@ -102,11 +102,11 @@ func (ims *HnswVectorIndex) Search(ctx context.Context, query []float32, maxResu
 
 func (ims *HnswVectorIndex) SearchWithKey(ctx context.Context, queryKey string, maxResults int, filter index.SearchFilter) (utils.MaxTupleHeap, error) {
 	ims.mu.RLock()
+	defer ims.mu.RUnlock()
 	query, found := ims.HnswIndex.Lookup(queryKey)
 	if !found {
 		return nil, fmt.Errorf("key not found")
 	}
-	ims.mu.RUnlock()
 	if query == nil {
 		return nil, nil
 	}

runtime/collections/in_mem/sequential/vector_index.go

@@ -103,8 +103,8 @@ func (ims *SequentialVectorIndex) Search(ctx context.Context, query []float32, m
 
 func (ims *SequentialVectorIndex) SearchWithKey(ctx context.Context, queryKey string, maxResults int, filter index.SearchFilter) (utils.MaxTupleHeap, error) {
 	ims.mu.RLock()
+	defer ims.mu.RUnlock()
 	query := ims.VectorMap[queryKey]
-	ims.mu.RUnlock()
 	if query == nil {
 		return nil, nil
 	}

runtime/collections/vector.go

@@ -98,7 +98,7 @@ func deleteIndexesNotInManifest(ctx context.Context, man *manifest.Manifest) {
 				Msg("Failed to find collection.")
 			continue
 		}
-		for _, collNs := range col.collectionNamespaceMap {
+		for _, collNs := range col.getCollectionNamespaceMap() {
 			vectorIndexMap := collNs.GetVectorIndexMap()
 			if vectorIndexMap == nil {
 				continue
@@ -160,7 +160,7 @@ func processManifestCollections(ctx context.Context, man *manifest.Manifest) {
 				}
 			}
 		}
-		for _, collNs := range col.collectionNamespaceMap {
+		for _, collNs := range col.getCollectionNamespaceMap() {
 			for searchMethodName, searchMethod := range collectionInfo.SearchMethods {
 				vi, err := collNs.GetVectorIndex(ctx, searchMethodName)
 

runtime/db/db.go

@@ -53,7 +53,7 @@ type runtimePostgresWriter struct {
 	buffer chan inferenceHistory
 	quit   chan struct{}
 	done   chan struct{}
-	mu     sync.RWMutex
+	once   sync.Once
 }
 
 type inferenceHistory struct {
@@ -67,37 +67,39 @@ type inferenceHistory struct {
 }
 
 func (w *runtimePostgresWriter) GetPool(ctx context.Context) (*pgxpool.Pool, error) {
-	w.mu.RLock()
+	var initErr error
+	w.once.Do(func() {
+		var connStr string
+		var err error
+		if secrets.HasSecret("MODUS_DB") {
+			connStr, err = secrets.GetSecretValue("MODUS_DB")
+		} else if secrets.HasSecret("HYPERMODE_METADATA_DB") {
+			// fallback to old secret name
+			// TODO: remove this after the transition is complete
+			connStr, err = secrets.GetSecretValue("HYPERMODE_METADATA_DB")
+		} else {
+			return
+		}
+
+		if err != nil {
+			initErr = err
+			return
+		}
+
+		if pool, err := pgxpool.New(ctx, connStr); err != nil {
+			initErr = err
+		} else {
+			w.dbpool = pool
+		}
+	})
+
 	if w.dbpool != nil {
-		defer w.mu.RUnlock()
 		return w.dbpool, nil
-	}
-	w.mu.RUnlock()
-
-	w.mu.Lock()
-	defer w.mu.Unlock()
-
-	var connStr string
-	var err error
-	if secrets.HasSecret("MODUS_DB") {
-		connStr, err = secrets.GetSecretValue("MODUS_DB")
-	} else if secrets.HasSecret("HYPERMODE_METADATA_DB") {
-		// fallback to old secret name
-		// TODO: remove this after the transition is complete
-		connStr, err = secrets.GetSecretValue("HYPERMODE_METADATA_DB")
+	} else if initErr != nil {
+		return nil, initErr
 	} else {
 		return nil, errDbNotConfigured
 	}
-	if err != nil {
-		return nil, err
-	}
-
-	if pool, err := pgxpool.New(ctx, connStr); err != nil {
-		return nil, err
-	} else {
-		w.dbpool = pool
-		return pool, nil
-	}
 }
 
 func (w *runtimePostgresWriter) Write(data inferenceHistory) {

runtime/dgraphclient/registry.go

@@ -14,7 +14,6 @@ import (
 	"crypto/x509"
 	"fmt"
 	"strings"
-	"sync"
 
 	"github.com/hypermodeinc/modus/lib/manifest"
 	"github.com/hypermodeinc/modus/runtime/manifestdata"
@@ -26,13 +25,13 @@ import (
 
 	"github.com/dgraph-io/dgo/v240"
 	"github.com/dgraph-io/dgo/v240/protos/api"
+	"github.com/puzpuzpuz/xsync/v3"
 )
 
 var dgr = newDgraphRegistry()
 
 type dgraphRegistry struct {
-	sync.RWMutex
-	dgraphConnectorCache map[string]*dgraphConnector
+	cache *xsync.MapOf[string, *dgraphConnector]
 }
 
 type authCreds struct {
@@ -51,35 +50,41 @@ func (a *authCreds) RequireTransportSecurity() bool {
 
 func newDgraphRegistry() *dgraphRegistry {
 	return &dgraphRegistry{
-		dgraphConnectorCache: make(map[string]*dgraphConnector),
+		cache: xsync.NewMapOf[string, *dgraphConnector](),
 	}
 }
 
 func ShutdownConns() {
-	dgr.Lock()
-	defer dgr.Unlock()
-	for _, ds := range dgr.dgraphConnectorCache {
-		ds.conn.Close()
-	}
-	clear(dgr.dgraphConnectorCache)
+	dgr.cache.Range(func(key string, _ *dgraphConnector) bool {
+		if connector, ok := dgr.cache.LoadAndDelete(key); ok {
+			connector.conn.Close()
+		}
+		return true
+	})
 }
 
 func (dr *dgraphRegistry) getDgraphConnector(ctx context.Context, dgName string) (*dgraphConnector, error) {
-	dr.RLock()
-	ds, ok := dr.dgraphConnectorCache[dgName]
-	dr.RUnlock()
-	if ok {
-		return ds, nil
-	}
-
-	dr.Lock()
-	defer dr.Unlock()
+	var creationErr error
+	ds, _ := dr.cache.LoadOrCompute(dgName, func() *dgraphConnector {
+		conn, err := createConnector(ctx, dgName)
+		if err != nil {
+			creationErr = err
+			return nil
+		}
+		return conn
+	})
 
-	if ds, ok := dr.dgraphConnectorCache[dgName]; ok {
-		return ds, nil
+	if creationErr != nil {
+		dr.cache.Delete(dgName)
+		return nil, creationErr
 	}
 
-	info, ok := manifestdata.GetManifest().Connections[dgName]
+	return ds, nil
+}
+
+func createConnector(ctx context.Context, dgName string) (*dgraphConnector, error) {
+	man := manifestdata.GetManifest()
+	info, ok := man.Connections[dgName]
 	if !ok {
 		return nil, fmt.Errorf("dgraph connection [%s] not found", dgName)
 	}
@@ -130,10 +135,10 @@ func (dr *dgraphRegistry) getDgraphConnector(ctx context.Context, dgName string)
 		return nil, err
 	}
 
-	ds = &dgraphConnector{
+	ds := &dgraphConnector{
 		conn:     conn,
 		dgClient: dgo.NewDgraphClient(api.NewDgraphClient(conn)),
 	}
-	dr.dgraphConnectorCache[dgName] = ds
+
 	return ds, nil
 }

runtime/go.mod

@@ -35,6 +35,7 @@ require (
 	github.com/neo4j/neo4j-go-driver/v5 v5.27.0
 	github.com/prometheus/client_golang v1.20.5
 	github.com/prometheus/common v0.60.1
+	github.com/puzpuzpuz/xsync/v3 v3.4.0
 	github.com/rs/cors v1.11.1
 	github.com/rs/xid v1.6.0
 	github.com/rs/zerolog v1.33.0

runtime/go.sum

@@ -230,6 +230,8 @@ github.com/prometheus/common v0.60.1 h1:FUas6GcOw66yB/73KC+BOZoFJmbo/1pojoILArPA
 github.com/prometheus/common v0.60.1/go.mod h1:h0LYf1R1deLSKtD4Vdg8gy4RuOvENW2J/h19V5NADQw=
 github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
 github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
+github.com/puzpuzpuz/xsync/v3 v3.4.0 h1:DuVBAdXuGFHv8adVXjWWZ63pJq+NRXOWVXlKDBZ+mJ4=
+github.com/puzpuzpuz/xsync/v3 v3.4.0/go.mod h1:VjzYrABPabuM4KyBh1Ftq6u8nhwY5tBPKP9jpmh0nnA=
 github.com/r3labs/sse/v2 v2.10.0 h1:hFEkLLFY4LDifoHdiCN/LlGBAdVJYsANaLqNYa1l/v0=
 github.com/r3labs/sse/v2 v2.10.0/go.mod h1:Igau6Whc+F17QUgML1fYe1VPZzTV6EMCnYktEmkNJ7I=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=