fix: cache tenant routing on message hot path

server/db/mysql/adapter.go

@@ -24,6 +24,7 @@ import (
 	immetrics "github.com/tinode/chat/server/metrics"
 	"github.com/tinode/chat/server/store"
 	t "github.com/tinode/chat/server/store/types"
+	"golang.org/x/sync/singleflight"
 )
 
 // adapter holds MySQL connection data.
@@ -43,14 +44,17 @@ type adapter struct {
 	// DB transaction timeout.
 	txTimeout time.Duration
 
-	baseConfig       *configType
-	tenantDBRouting  bool
-	tenantScopedDB   bool
-	tenantID         t.TenantID
-	tenantRegistryDB string
-	tenantAdapters   map[t.TenantID]*adapter
-	tenantMu         sync.Mutex
-	groupShardWindow string
+	baseConfig         *configType
+	tenantDBRouting    bool
+	tenantScopedDB     bool
+	tenantID           t.TenantID
+	tenantRegistryDB   string
+	tenantAdapters     map[t.TenantID]*adapter
+	tenantRouteExpiry  map[t.TenantID]time.Time
+	tenantRouteGen     uint64
+	tenantRouteRefresh singleflight.Group
+	tenantMu           sync.Mutex
+	groupShardWindow   string
 }
 
 const (
@@ -63,6 +67,7 @@ const (
 	defaultMaxResults = 1024
 	// This is capped by the Session's send queue limit (128).
 	defaultMaxMessageResults = 100
+	tenantRouteCacheTTL      = time.Minute
 
 	// If DB request timeout is specified,
 	// we allocate txTimeoutMultiplier times more time for transactions.
@@ -262,6 +267,7 @@ func (a *adapter) Open(jsonconfig json.RawMessage) error {
 	a.tenantID = defaultTenantID
 	a.tenantRegistryDB = strings.TrimSpace(config.TenantRegistryDB)
 	a.tenantAdapters = make(map[t.TenantID]*adapter)
+	a.tenantRouteExpiry = make(map[t.TenantID]time.Time)
 
 	if a.maxResults <= 0 {
 		a.maxResults = defaultMaxResults
@@ -320,6 +326,8 @@ func (a *adapter) Close() error {
 		}
 		delete(a.tenantAdapters, tenantID)
 	}
+	a.tenantRouteExpiry = make(map[t.TenantID]time.Time)
+	a.tenantRouteGen++
 	a.tenantMu.Unlock()
 	if a.registryDB != nil {
 		if closeErr := a.registryDB.Close(); err == nil {
@@ -475,35 +483,62 @@ func (a *adapter) TenantListActive(limit int) ([]t.Tenant, error) {
 }
 
 func (a *adapter) TenantAdapterFor(tenantID t.TenantID) (dbadapter.TenantBusinessAdapter, dbadapter.TenantAuthAdapter, error) {
-	if !a.tenantDBRouting {
-		if a.tenantID != defaultTenantID && a.tenantID != tenantID {
-			return nil, nil, t.ErrMalformed
-		}
-		if a.registryDB != nil {
-			tenant, err := a.tenantByID(tenantID)
-			if err != nil {
-				return nil, nil, err
-			}
-			if tenant.State != t.TenantStateActive {
-				return nil, nil, t.ErrPermissionDenied
-			}
-		}
-		return a, a, nil
+	if !tenantID.IsValid() || (!a.tenantDBRouting && a.tenantID != defaultTenantID && a.tenantID != tenantID) {
+		return nil, nil, t.ErrMalformed
 	}
-	tenant, err := a.tenantByID(tenantID)
+	if tenantAdp, ok := a.cachedTenantAdapter(tenantID, time.Now()); ok {
+		return tenantAdp, tenantAdp, nil
+	}
+
+	resolved, err, _ := a.tenantRouteRefresh.Do(strconv.FormatInt(int64(tenantID), 10), func() (any, error) {
+		if tenantAdp, ok := a.cachedTenantAdapter(tenantID, time.Now()); ok {
+			return tenantAdp, nil
+		}
+		generation := a.tenantRouteGeneration()
+		if a.registryDB != nil {
+			if _, err := a.tenantByID(tenantID); err != nil {
+				return nil, err
+			}
+		}
+		tenantAdp := a.tenantAdapter(tenantID)
+		a.cacheTenantAdapter(tenantID, tenantAdp, generation, time.Now().Add(tenantRouteCacheTTL))
+		return tenantAdp, nil
+	})
 	if err != nil {
 		return nil, nil, err
 	}
-	if tenant.State != t.TenantStateActive {
-		return nil, nil, t.ErrPermissionDenied
+	tenantAdp := resolved.(*adapter)
+	return tenantAdp, tenantAdp, nil
+}
+
+func (a *adapter) cachedTenantAdapter(tenantID t.TenantID, now time.Time) (*adapter, bool) {
+	a.tenantMu.Lock()
+	defer a.tenantMu.Unlock()
+	if expiresAt, ok := a.tenantRouteExpiry[tenantID]; !ok || !now.Before(expiresAt) {
+		return nil, false
+	}
+	if !a.tenantDBRouting {
+		return a, true
+	}
+	tenantAdp := a.tenantAdapters[tenantID]
+	return tenantAdp, tenantAdp != nil
+}
+
+func (a *adapter) tenantRouteGeneration() uint64 {
+	a.tenantMu.Lock()
+	defer a.tenantMu.Unlock()
+	return a.tenantRouteGen
+}
+
+func (a *adapter) tenantAdapter(tenantID t.TenantID) *adapter {
+	if !a.tenantDBRouting {
+		return a
 	}
 	a.tenantMu.Lock()
-	if tenantAdp := a.tenantAdapters[tenantID]; tenantAdp != nil {
-		a.tenantMu.Unlock()
-		return tenantAdp, tenantAdp, nil
+	defer a.tenantMu.Unlock()
+	if existing := a.tenantAdapters[tenantID]; existing != nil {
+		return existing
 	}
-	a.tenantMu.Unlock()
-
 	registryDB := a.registryDB
 	if registryDB == nil {
 		registryDB = a.db
@@ -525,14 +560,32 @@ func (a *adapter) TenantAdapterFor(tenantID t.TenantID) (dbadapter.TenantBusines
 		// explicit tenant_id writes enabled for logical isolation.
 		tenantScopedDB: false,
 	}
-	a.tenantMu.Lock()
-	if existing := a.tenantAdapters[tenantID]; existing != nil {
-		a.tenantMu.Unlock()
-		return existing, existing, nil
-	}
 	a.tenantAdapters[tenantID] = tenantAdp
-	a.tenantMu.Unlock()
-	return tenantAdp, tenantAdp, nil
+	return tenantAdp
+}
+
+func (a *adapter) cacheTenantAdapter(tenantID t.TenantID, tenantAdp *adapter, generation uint64, expiresAt time.Time) {
+	a.tenantMu.Lock()
+	defer a.tenantMu.Unlock()
+	if generation != a.tenantRouteGen || tenantAdp == nil {
+		return
+	}
+	if a.tenantRouteExpiry == nil {
+		a.tenantRouteExpiry = make(map[t.TenantID]time.Time)
+	}
+	a.tenantRouteExpiry[tenantID] = expiresAt
+}
+
+// InvalidateTenantRoute forces the next operation to revalidate tenant state.
+func (a *adapter) InvalidateTenantRoute(tenantID t.TenantID) {
+	a.tenantMu.Lock()
+	defer a.tenantMu.Unlock()
+	a.tenantRouteGen++
+	if tenantID.IsValid() {
+		delete(a.tenantRouteExpiry, tenantID)
+		return
+	}
+	a.tenantRouteExpiry = make(map[t.TenantID]time.Time)
 }
 
 func (a *adapter) tenantByID(tenantID t.TenantID) (*t.Tenant, error) {