SHOW tsa_cache (#1529)

Author: Denchick

docs/sharding/console/sql_commands.mdx

@@ -149,7 +149,7 @@ This command is used to retrieve information about a specific entity of SPQR clu
 ```sql
 SHOW entity [WHERE where_clause]
 
-where entity is one of databases, routers, pools, shards, backend_connections, key_ranges, clients, status, distributions, version, relations, sequences, task_groups, prepared_statements:
+where entity is one of backend_connections, clients, databases, distributions, key_ranges, pools, prepared_statements, relations, routers, sequences, shards, status, task_groups, tsa_cache, users, version:
 ```
 
 The`where_clause` is used to filter the results based on specific conditions. It allows you to specify criteria for selecting the information to be shown.

pkg/clientinteractor/interactor.go

@@ -1817,6 +1817,33 @@ func (pi *PSQLInteractor) Users(ctx context.Context) error {
 	return pi.CompleteMsg(len(berules))
 }
 
+// TsaCache outputs TSA cache entries showing the target session attributes cache status
+func (pi *PSQLInteractor) TsaCache(ctx context.Context, cacheEntries map[pool.TsaKey]pool.CachedEntry) error {
+	if err := pi.WriteHeader("tsa", "host", "az", "alive", "match", "reason", "last_check_time"); err != nil {
+		spqrlog.Zero.Error().Err(err).Msg("")
+		return err
+	}
+
+	count := 0
+	for key, entry := range cacheEntries {
+		if err := pi.WriteDataRow(
+			string(key.Tsa),
+			key.Host,
+			key.AZ,
+			fmt.Sprintf("%v", entry.Result.Alive),
+			fmt.Sprintf("%v", entry.Result.Match),
+			entry.Result.Reason,
+			entry.LastCheckTime.Format(time.RFC3339),
+		); err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("")
+			return err
+		}
+		count++
+	}
+
+	return pi.CompleteMsg(count)
+}
+
 // Outputs groupBy get list values and counts its 'groupByCol' property.
 // 'groupByCol' sorted in grouped result by string key ASC mode
 //

pkg/connmgr/connmgr.go

@@ -1,13 +1,15 @@
 package connmgr
 
 import (
+	"github.com/pg-sharding/spqr/pkg/pool"
 	"github.com/pg-sharding/spqr/pkg/tsa"
 )
 
 type ConnectionStatsMgr interface {
 	ConnectionIterator
 
 	InstanceHealthChecks() map[string]tsa.CachedCheckResult
+	TsaCacheEntries() map[pool.TsaKey]pool.CachedEntry
 
 	/*
 		user-facing connection stat callbacks.

pkg/coord/clustered_coord.go

@@ -59,6 +59,11 @@ func (ci grpcConnMgr) InstanceHealthChecks() map[string]tsa.CachedCheckResult {
 	return map[string]tsa.CachedCheckResult{}
 }
 
+// TsaCacheEntries implements connmgr.ConnectionStatsMgr.
+func (ci grpcConnMgr) TsaCacheEntries() map[pool.TsaKey]pool.CachedEntry {
+	return map[pool.TsaKey]pool.CachedEntry{}
+}
+
 // TODO implement it
 // ActiveTcpCount implements connectiterator.ConnectIterator.
 func (ci grpcConnMgr) ActiveTcpCount() int64 {

pkg/meta/meta.go

@@ -880,6 +880,9 @@ func ProcessShow(ctx context.Context, stmt *spqrparser.Show, mngr EntityMgr, ci
 		return cli.MoveStats(ctx, stats)
 	case spqrparser.Users:
 		return cli.Users(ctx)
+	case spqrparser.TsaCacheStr:
+		cacheEntries := ci.TsaCacheEntries()
+		return cli.TsaCache(ctx, cacheEntries)
 
 	default:
 		return ErrUnknownCoordinatorCommand

pkg/pool/dbpool.go

@@ -46,6 +46,11 @@ func (s *DBPool) InstanceHealthChecks() map[string]tsa.CachedCheckResult {
 	return s.checker.InstanceHealthChecks()
 }
 
+// TsaCacheEntries implements MultiShardTSAPool.
+func (s *DBPool) TsaCacheEntries() map[TsaKey]CachedEntry {
+	return s.cache.GetAllCachedEntries()
+}
+
 // View implements MultiShardPool.
 func (s *DBPool) View() Statistics {
 	return s.pool.View()

pkg/pool/dbpool_cache.go

@@ -155,6 +155,20 @@ func (c *DbpoolCache) GetAllEntries() map[TsaKey]LocalCheckResult {
 	return entries
 }
 
+// GetAllCachedEntries returns all cache entries with timestamps
+func (c *DbpoolCache) GetAllCachedEntries() map[TsaKey]CachedEntry {
+	entries := make(map[TsaKey]CachedEntry)
+	c.cache.Range(func(key, value any) bool {
+		if tsaKey, ok := key.(TsaKey); ok {
+			if entry, ok := value.(CachedEntry); ok {
+				entries[tsaKey] = entry
+			}
+		}
+		return true
+	})
+	return entries
+}
+
 // startCacheCleanup starts a background goroutine that cleans up stale cache entries every 30 seconds
 func (c *DbpoolCache) startCacheCleanup(d time.Duration) {
 	go func() {

pkg/pool/pool.go

@@ -60,6 +60,7 @@ type MultiShardTSAPool interface {
 	ShardMapping() map[string]*config.Shard
 	ConnectionWithTSA(clid uint, key kr.ShardKey, targetSessionAttrs tsa.TSA) (shard.ShardHostInstance, error)
 	InstanceHealthChecks() map[string]tsa.CachedCheckResult
+	TsaCacheEntries() map[TsaKey]CachedEntry
 
 	StopCacheWatchdog()
 }

router/rulerouter/rulerouter.go

@@ -11,6 +11,7 @@ import (
 	"github.com/pg-sharding/spqr/pkg/catalog"
 	"github.com/pg-sharding/spqr/pkg/connmgr"
 	"github.com/pg-sharding/spqr/pkg/models/spqrerror"
+	"github.com/pg-sharding/spqr/pkg/pool"
 	"github.com/pg-sharding/spqr/pkg/tsa"
 	"github.com/pg-sharding/spqr/pkg/txstatus"
 	"golang.org/x/sync/semaphore"
@@ -82,6 +83,19 @@ func (r *RuleRouterImpl) InstanceHealthChecks() map[string]tsa.CachedCheckResult
 	return rt
 }
 
+// TsaCacheEntries implements ConnectionStatsMgr.
+func (r *RuleRouterImpl) TsaCacheEntries() map[pool.TsaKey]pool.CachedEntry {
+	rt := map[pool.TsaKey]pool.CachedEntry{}
+	_ = r.NotifyRoutes(func(r *route.Route) (bool, error) {
+		m := r.ServPool().TsaCacheEntries()
+		for k, v := range m {
+			rt[k] = v
+		}
+		return true, nil
+	})
+	return rt
+}
+
 // ReleaseConnection implements RuleRouter.
 func (r *RuleRouterImpl) ReleaseConnection() {
 	r.activeTcpCount.Add(-1)

test/regress/schedule/console

@@ -19,4 +19,5 @@ test: sequence
 test: alter_relation
 test: reference_table
 test: show_users
+test: show_tsa_cache
 test: routing_expression