add custom store cache

management/server/store/cache/dual_key_cache.go

@@ -0,0 +1,129 @@
+package cache
+
+import (
+	"context"
+	"sync"
+)
+
+// DualKeyCache provides a caching mechanism where each entry has two keys:
+// - Primary key (e.g., objectID): used for accessing and invalidating specific entries
+// - Secondary key (e.g., accountID): used for bulk invalidation of all entries with the same secondary key
+type DualKeyCache[K1 comparable, K2 comparable, V any] struct {
+	mu             sync.RWMutex
+	primaryIndex   map[K1]V               // Primary key -> Value
+	secondaryIndex map[K2]map[K1]struct{} // Secondary key -> Set of primary keys
+	reverseLookup  map[K1]K2              // Primary key -> Secondary key
+}
+
+// NewDualKeyCache creates a new dual-key cache
+func NewDualKeyCache[K1 comparable, K2 comparable, V any]() *DualKeyCache[K1, K2, V] {
+	return &DualKeyCache[K1, K2, V]{
+		primaryIndex:   make(map[K1]V),
+		secondaryIndex: make(map[K2]map[K1]struct{}),
+		reverseLookup:  make(map[K1]K2),
+	}
+}
+
+// Get retrieves a value from the cache using the primary key
+func (c *DualKeyCache[K1, K2, V]) Get(ctx context.Context, primaryKey K1) (V, bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	value, ok := c.primaryIndex[primaryKey]
+	return value, ok
+}
+
+// Set stores a value in the cache with both primary and secondary keys
+func (c *DualKeyCache[K1, K2, V]) Set(ctx context.Context, primaryKey K1, secondaryKey K2, value V) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if oldSecondaryKey, exists := c.reverseLookup[primaryKey]; exists {
+		if primaryKeys, ok := c.secondaryIndex[oldSecondaryKey]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.secondaryIndex, oldSecondaryKey)
+			}
+		}
+	}
+
+	c.primaryIndex[primaryKey] = value
+	c.reverseLookup[primaryKey] = secondaryKey
+
+	if _, exists := c.secondaryIndex[secondaryKey]; !exists {
+		c.secondaryIndex[secondaryKey] = make(map[K1]struct{})
+	}
+	c.secondaryIndex[secondaryKey][primaryKey] = struct{}{}
+}
+
+// InvalidateByPrimaryKey removes an entry using the primary key
+func (c *DualKeyCache[K1, K2, V]) InvalidateByPrimaryKey(ctx context.Context, primaryKey K1) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if secondaryKey, exists := c.reverseLookup[primaryKey]; exists {
+		if primaryKeys, ok := c.secondaryIndex[secondaryKey]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.secondaryIndex, secondaryKey)
+			}
+		}
+		delete(c.reverseLookup, primaryKey)
+	}
+
+	delete(c.primaryIndex, primaryKey)
+}
+
+// InvalidateBySecondaryKey removes all entries with the given secondary key
+func (c *DualKeyCache[K1, K2, V]) InvalidateBySecondaryKey(ctx context.Context, secondaryKey K2) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	primaryKeys, exists := c.secondaryIndex[secondaryKey]
+	if !exists {
+		return
+	}
+
+	for primaryKey := range primaryKeys {
+		delete(c.primaryIndex, primaryKey)
+		delete(c.reverseLookup, primaryKey)
+	}
+
+	delete(c.secondaryIndex, secondaryKey)
+}
+
+// InvalidateAll removes all entries from the cache
+func (c *DualKeyCache[K1, K2, V]) InvalidateAll(ctx context.Context) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.primaryIndex = make(map[K1]V)
+	c.secondaryIndex = make(map[K2]map[K1]struct{})
+	c.reverseLookup = make(map[K1]K2)
+}
+
+// Size returns the number of entries in the cache
+func (c *DualKeyCache[K1, K2, V]) Size() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return len(c.primaryIndex)
+}
+
+// GetOrSet retrieves a value from the cache, or sets it using the provided function if not found
+// The loadFunc should return both the value and the secondary key (extracted from the value)
+func (c *DualKeyCache[K1, K2, V]) GetOrSet(ctx context.Context, primaryKey K1, loadFunc func() (V, K2, error)) (V, error) {
+	if value, ok := c.Get(ctx, primaryKey); ok {
+		return value, nil
+	}
+
+	value, secondaryKey, err := loadFunc()
+	if err != nil {
+		var zero V
+		return zero, err
+	}
+
+	c.Set(ctx, primaryKey, secondaryKey, value)
+
+	return value, nil
+}

management/server/store/cache/single_key_cache.go

@@ -0,0 +1,77 @@
+package cache
+
+import (
+	"context"
+	"sync"
+)
+
+// SingleKeyCache provides a simple caching mechanism with a single key
+type SingleKeyCache[K comparable, V any] struct {
+	mu    sync.RWMutex
+	cache map[K]V // Key -> Value
+}
+
+// NewSingleKeyCache creates a new single-key cache
+func NewSingleKeyCache[K comparable, V any]() *SingleKeyCache[K, V] {
+	return &SingleKeyCache[K, V]{
+		cache: make(map[K]V),
+	}
+}
+
+// Get retrieves a value from the cache using the key
+func (c *SingleKeyCache[K, V]) Get(ctx context.Context, key K) (V, bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	value, ok := c.cache[key]
+	return value, ok
+}
+
+// Set stores a value in the cache with the given key
+func (c *SingleKeyCache[K, V]) Set(ctx context.Context, key K, value V) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.cache[key] = value
+}
+
+// Invalidate removes an entry using the key
+func (c *SingleKeyCache[K, V]) Invalidate(ctx context.Context, key K) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	delete(c.cache, key)
+}
+
+// InvalidateAll removes all entries from the cache
+func (c *SingleKeyCache[K, V]) InvalidateAll(ctx context.Context) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.cache = make(map[K]V)
+}
+
+// Size returns the number of entries in the cache
+func (c *SingleKeyCache[K, V]) Size() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return len(c.cache)
+}
+
+// GetOrSet retrieves a value from the cache, or sets it using the provided function if not found
+func (c *SingleKeyCache[K, V]) GetOrSet(ctx context.Context, key K, loadFunc func() (V, error)) (V, error) {
+	if value, ok := c.Get(ctx, key); ok {
+		return value, nil
+	}
+
+	value, err := loadFunc()
+	if err != nil {
+		var zero V
+		return zero, err
+	}
+
+	c.Set(ctx, key, value)
+
+	return value, nil
+}

management/server/store/cache/triple_key_cache.go

@@ -0,0 +1,242 @@
+package cache
+
+import (
+	"context"
+	"sync"
+)
+
+// TripleKeyCache provides a caching mechanism where each entry has three keys:
+// - Primary key (K1): used for accessing and invalidating specific entries
+// - Secondary key (K2): used for bulk invalidation of all entries with the same secondary key
+// - Tertiary key (K3): used for bulk invalidation of all entries with the same tertiary key
+type TripleKeyCache[K1 comparable, K2 comparable, K3 comparable, V any] struct {
+	mu             sync.RWMutex
+	primaryIndex   map[K1]V               // Primary key -> Value
+	secondaryIndex map[K2]map[K1]struct{} // Secondary key -> Set of primary keys
+	tertiaryIndex  map[K3]map[K1]struct{} // Tertiary key -> Set of primary keys
+	reverseLookup  map[K1]keyPair[K2, K3] // Primary key -> Secondary and Tertiary keys
+}
+
+type keyPair[K2 comparable, K3 comparable] struct {
+	secondary K2
+	tertiary  K3
+}
+
+// NewTripleKeyCache creates a new triple-key cache
+func NewTripleKeyCache[K1 comparable, K2 comparable, K3 comparable, V any]() *TripleKeyCache[K1, K2, K3, V] {
+	return &TripleKeyCache[K1, K2, K3, V]{
+		primaryIndex:   make(map[K1]V),
+		secondaryIndex: make(map[K2]map[K1]struct{}),
+		tertiaryIndex:  make(map[K3]map[K1]struct{}),
+		reverseLookup:  make(map[K1]keyPair[K2, K3]),
+	}
+}
+
+// Get retrieves a value from the cache using the primary key
+func (c *TripleKeyCache[K1, K2, K3, V]) Get(ctx context.Context, primaryKey K1) (V, bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	value, ok := c.primaryIndex[primaryKey]
+	return value, ok
+}
+
+// Set stores a value in the cache with primary, secondary, and tertiary keys
+func (c *TripleKeyCache[K1, K2, K3, V]) Set(ctx context.Context, primaryKey K1, secondaryKey K2, tertiaryKey K3, value V) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if oldKeys, exists := c.reverseLookup[primaryKey]; exists {
+		if primaryKeys, ok := c.secondaryIndex[oldKeys.secondary]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.secondaryIndex, oldKeys.secondary)
+			}
+		}
+		if primaryKeys, ok := c.tertiaryIndex[oldKeys.tertiary]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.tertiaryIndex, oldKeys.tertiary)
+			}
+		}
+	}
+
+	c.primaryIndex[primaryKey] = value
+	c.reverseLookup[primaryKey] = keyPair[K2, K3]{
+		secondary: secondaryKey,
+		tertiary:  tertiaryKey,
+	}
+
+	if _, exists := c.secondaryIndex[secondaryKey]; !exists {
+		c.secondaryIndex[secondaryKey] = make(map[K1]struct{})
+	}
+	c.secondaryIndex[secondaryKey][primaryKey] = struct{}{}
+
+	if _, exists := c.tertiaryIndex[tertiaryKey]; !exists {
+		c.tertiaryIndex[tertiaryKey] = make(map[K1]struct{})
+	}
+	c.tertiaryIndex[tertiaryKey][primaryKey] = struct{}{}
+}
+
+// InvalidateByPrimaryKey removes an entry using the primary key
+func (c *TripleKeyCache[K1, K2, K3, V]) InvalidateByPrimaryKey(ctx context.Context, primaryKey K1) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if keys, exists := c.reverseLookup[primaryKey]; exists {
+		if primaryKeys, ok := c.secondaryIndex[keys.secondary]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.secondaryIndex, keys.secondary)
+			}
+		}
+		if primaryKeys, ok := c.tertiaryIndex[keys.tertiary]; ok {
+			delete(primaryKeys, primaryKey)
+			if len(primaryKeys) == 0 {
+				delete(c.tertiaryIndex, keys.tertiary)
+			}
+		}
+		delete(c.reverseLookup, primaryKey)
+	}
+
+	delete(c.primaryIndex, primaryKey)
+}
+
+// InvalidateBySecondaryKey removes all entries with the given secondary key
+func (c *TripleKeyCache[K1, K2, K3, V]) InvalidateBySecondaryKey(ctx context.Context, secondaryKey K2) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	primaryKeys, exists := c.secondaryIndex[secondaryKey]
+	if !exists {
+		return
+	}
+
+	for primaryKey := range primaryKeys {
+		if keys, ok := c.reverseLookup[primaryKey]; ok {
+			if tertiaryPrimaryKeys, exists := c.tertiaryIndex[keys.tertiary]; exists {
+				delete(tertiaryPrimaryKeys, primaryKey)
+				if len(tertiaryPrimaryKeys) == 0 {
+					delete(c.tertiaryIndex, keys.tertiary)
+				}
+			}
+		}
+		delete(c.primaryIndex, primaryKey)
+		delete(c.reverseLookup, primaryKey)
+	}
+
+	delete(c.secondaryIndex, secondaryKey)
+}
+
+// InvalidateByTertiaryKey removes all entries with the given tertiary key
+func (c *TripleKeyCache[K1, K2, K3, V]) InvalidateByTertiaryKey(ctx context.Context, tertiaryKey K3) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	primaryKeys, exists := c.tertiaryIndex[tertiaryKey]
+	if !exists {
+		return
+	}
+
+	for primaryKey := range primaryKeys {
+		if keys, ok := c.reverseLookup[primaryKey]; ok {
+			if secondaryPrimaryKeys, exists := c.secondaryIndex[keys.secondary]; exists {
+				delete(secondaryPrimaryKeys, primaryKey)
+				if len(secondaryPrimaryKeys) == 0 {
+					delete(c.secondaryIndex, keys.secondary)
+				}
+			}
+		}
+		delete(c.primaryIndex, primaryKey)
+		delete(c.reverseLookup, primaryKey)
+	}
+
+	delete(c.tertiaryIndex, tertiaryKey)
+}
+
+// InvalidateAll removes all entries from the cache
+func (c *TripleKeyCache[K1, K2, K3, V]) InvalidateAll(ctx context.Context) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.primaryIndex = make(map[K1]V)
+	c.secondaryIndex = make(map[K2]map[K1]struct{})
+	c.tertiaryIndex = make(map[K3]map[K1]struct{})
+	c.reverseLookup = make(map[K1]keyPair[K2, K3])
+}
+
+// Size returns the number of entries in the cache
+func (c *TripleKeyCache[K1, K2, K3, V]) Size() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return len(c.primaryIndex)
+}
+
+// GetOrSet retrieves a value from the cache, or sets it using the provided function if not found
+// The loadFunc should return the value, secondary key, and tertiary key (extracted from the value)
+func (c *TripleKeyCache[K1, K2, K3, V]) GetOrSet(ctx context.Context, primaryKey K1, loadFunc func() (V, K2, K3, error)) (V, error) {
+	if value, ok := c.Get(ctx, primaryKey); ok {
+		return value, nil
+	}
+
+	value, secondaryKey, tertiaryKey, err := loadFunc()
+	if err != nil {
+		var zero V
+		return zero, err
+	}
+
+	c.Set(ctx, primaryKey, secondaryKey, tertiaryKey, value)
+
+	return value, nil
+}
+
+// GetOrSetBySecondaryKey retrieves a value from the cache using the secondary key, or sets it using the provided function if not found
+// The loadFunc should return the value, primary key, secondary key, and tertiary key
+func (c *TripleKeyCache[K1, K2, K3, V]) GetOrSetBySecondaryKey(ctx context.Context, secondaryKey K2, loadFunc func() (V, K1, K3, error)) (V, error) {
+	c.mu.RLock()
+	if primaryKeys, exists := c.secondaryIndex[secondaryKey]; exists && len(primaryKeys) > 0 {
+		for primaryKey := range primaryKeys {
+			if value, ok := c.primaryIndex[primaryKey]; ok {
+				c.mu.RUnlock()
+				return value, nil
+			}
+		}
+	}
+	c.mu.RUnlock()
+
+	value, primaryKey, tertiaryKey, err := loadFunc()
+	if err != nil {
+		var zero V
+		return zero, err
+	}
+
+	c.Set(ctx, primaryKey, secondaryKey, tertiaryKey, value)
+
+	return value, nil
+}
+
+// GetOrSetByTertiaryKey retrieves a value from the cache using the tertiary key, or sets it using the provided function if not found
+// The loadFunc should return the value, primary key, secondary key, and tertiary key
+func (c *TripleKeyCache[K1, K2, K3, V]) GetOrSetByTertiaryKey(ctx context.Context, tertiaryKey K3, loadFunc func() (V, K1, K2, error)) (V, error) {
+	c.mu.RLock()
+	if primaryKeys, exists := c.tertiaryIndex[tertiaryKey]; exists && len(primaryKeys) > 0 {
+		for primaryKey := range primaryKeys {
+			if value, ok := c.primaryIndex[primaryKey]; ok {
+				c.mu.RUnlock()
+				return value, nil
+			}
+		}
+	}
+	c.mu.RUnlock()
+
+	value, primaryKey, secondaryKey, err := loadFunc()
+	if err != nil {
+		var zero V
+		return zero, err
+	}
+
+	c.Set(ctx, primaryKey, secondaryKey, tertiaryKey, value)
+
+	return value, nil
+}

management/server/telemetry/store_metrics.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 )
 
@@ -14,6 +15,8 @@ type StoreMetrics struct {
 	persistenceDurationMicro           metric.Int64Histogram
 	persistenceDurationMs              metric.Int64Histogram
 	transactionDurationMs              metric.Int64Histogram
+	queryDurationMs                    metric.Int64Histogram
+	queryCounter                       metric.Int64Counter
 	ctx                                context.Context
 }
 
@@ -59,12 +62,29 @@ func NewStoreMetrics(ctx context.Context, meter metric.Meter) (*StoreMetrics, er
 		return nil, err
 	}
 
+	queryDurationMs, err := meter.Int64Histogram("management.store.query.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration of database query operations with operation type and table name"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	queryCounter, err := meter.Int64Counter("management.store.query.count",
+		metric.WithDescription("Count of database query operations with operation type, table name, and status"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
 	return &StoreMetrics{
 		globalLockAcquisitionDurationMicro: globalLockAcquisitionDurationMicro,
 		globalLockAcquisitionDurationMs:    globalLockAcquisitionDurationMs,
 		persistenceDurationMicro:           persistenceDurationMicro,
 		persistenceDurationMs:              persistenceDurationMs,
 		transactionDurationMs:              transactionDurationMs,
+		queryDurationMs:                    queryDurationMs,
+		queryCounter:                       queryCounter,
 		ctx:                                ctx,
 	}, nil
 }
@@ -85,3 +105,13 @@ func (metrics *StoreMetrics) CountPersistenceDuration(duration time.Duration) {
 func (metrics *StoreMetrics) CountTransactionDuration(duration time.Duration) {
 	metrics.transactionDurationMs.Record(metrics.ctx, duration.Milliseconds())
 }
+
+// CountStoreOperation records a store operation with its method name, status, and duration
+func (metrics *StoreMetrics) CountStoreOperation(method string, duration time.Duration) {
+	attrs := []attribute.KeyValue{
+		attribute.String("method", method),
+	}
+
+	metrics.queryDurationMs.Record(metrics.ctx, duration.Milliseconds(), metric.WithAttributes(attrs...))
+	metrics.queryCounter.Add(metrics.ctx, 1, metric.WithAttributes(attrs...))
+}