sql: reimplement sqlstats API using iterator

As we move to implement an aggregated virtual table of in-memory
sqlstats and persisted sqlstats, current SQL Stats API makes it
difficult to implement such behaviour.
This commit introduces lower level iterator APIs that can be used
later for the aggregated virtual table.

Release note: None

pkg/sql/sqlstats/sslocal/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "sql_stats.go",
         "sql_stats_controller.go",
+        "sslocal_iterator.go",
         "sslocal_provider.go",
         "sslocal_sink.go",
         "sslocal_stats_collector.go",

pkg/sql/sqlstats/sslocal/sql_stats.go

@@ -129,6 +129,7 @@ func (s *SQLStats) getStatsForApplication(appName string) *ssmemstorage.Containe
 		&s.atomic.uniqueStmtFingerprintCount,
 		&s.atomic.uniqueTxnFingerprintCount,
 		s.mu.mon,
+		appName,
 	)
 	s.mu.apps[appName] = a
 	return a

pkg/sql/sqlstats/sslocal/sslocal_iterator.go

@@ -0,0 +1,121 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package sslocal
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/ssmemstorage"
+)
+
+type baseIterator struct {
+	sqlStats *SQLStats
+	idx      int
+	appNames []string
+	options  *sqlstats.IteratorOptions
+}
+
+// StmtStatsIterator is an iterator that can be used to iterate over the
+// statement statistics stored in SQLStats.
+type StmtStatsIterator struct {
+	baseIterator
+	curIter *ssmemstorage.StmtStatsIterator
+}
+
+// NewStmtStatsIterator returns a new instance of the StmtStatsIterator.
+func NewStmtStatsIterator(s *SQLStats, options *sqlstats.IteratorOptions) *StmtStatsIterator {
+	appNames := s.getAppNames(options.SortedAppNames)
+
+	return &StmtStatsIterator{
+		baseIterator: baseIterator{
+			sqlStats: s,
+			idx:      -1,
+			appNames: appNames,
+			options:  options,
+		},
+	}
+}
+
+// Next increments the internal counter of the StmtStatsIterator. It returns
+// true if the following Cur() call will be valid, false otherwise.
+func (s *StmtStatsIterator) Next() bool {
+	// If we haven't called Next() for the first time or our current child
+	// iterator has finished iterator, then we increment s.idx.
+	if s.curIter == nil || !s.curIter.Next() {
+		s.idx++
+		if s.idx >= len(s.appNames) {
+			return false
+		}
+		statsContainer := s.sqlStats.getStatsForApplication(s.appNames[s.idx])
+		s.curIter = statsContainer.StmtStatsIterator(s.options)
+		return s.Next()
+	}
+
+	// This happens when our child iterator is valid and s.curIter.Next() call
+	// is true.
+	return true
+}
+
+// Cur returns the roachpb.CollectedStatementStatistics at the current internal
+// counter.
+func (s *StmtStatsIterator) Cur() *roachpb.CollectedStatementStatistics {
+	return s.curIter.Cur()
+}
+
+// TxnStatsIterator is an iterator that can be used to iterate over the
+// statement statistics stored in SQLStats.
+type TxnStatsIterator struct {
+	baseIterator
+	curIter *ssmemstorage.TxnStatsIterator
+}
+
+// NewTxnStatsIterator returns a new instance of the TxnStatsIterator.
+func NewTxnStatsIterator(s *SQLStats, options *sqlstats.IteratorOptions) *TxnStatsIterator {
+	appNames := s.getAppNames(options.SortedAppNames)
+
+	return &TxnStatsIterator{
+		baseIterator: baseIterator{
+			sqlStats: s,
+			idx:      -1,
+			appNames: appNames,
+			options:  options,
+		},
+	}
+}
+
+// Next increments the internal counter of the TxnStatsIterator. It returns
+// true if the following Cur() call will be valid, false otherwise.
+func (t *TxnStatsIterator) Next() bool {
+	// If we haven't called Next() for the first time or our current child
+	// iterator has finished iterator, then we increment s.idx.
+	if t.curIter == nil || !t.curIter.Next() {
+		t.idx++
+		if t.idx >= len(t.appNames) {
+			return false
+		}
+		statsContainer := t.sqlStats.getStatsForApplication(t.appNames[t.idx])
+		t.curIter = statsContainer.TxnStatsIterator(t.options)
+		return t.Next()
+	}
+
+	// This happens when our child iterator is valid and s.curIter.Next() call
+	// is true.
+	return true
+}
+
+// Cur returns the roachpb.CollectedTransactionStatistics at the current internal
+// counter.
+func (t *TxnStatsIterator) Cur() (
+	roachpb.TransactionFingerprintID,
+	*roachpb.CollectedTransactionStatistics,
+) {
+	return t.curIter.Cur()
+}

pkg/sql/sqlstats/sslocal/sslocal_provider.go

@@ -111,6 +111,7 @@ func (s *SQLStats) GetWriterForApplication(appName string) sqlstats.Writer {
 		&s.atomic.uniqueStmtFingerprintCount,
 		&s.atomic.uniqueTxnFingerprintCount,
 		s.mu.mon,
+		appName,
 	)
 	s.mu.apps[appName] = a
 	return a
@@ -127,36 +128,45 @@ func (s *SQLStats) GetLastReset() time.Time {
 func (s *SQLStats) IterateStatementStats(
 	ctx context.Context, options *sqlstats.IteratorOptions, visitor sqlstats.StatementVisitor,
 ) error {
-	appNames := s.getAppNames(options.SortedAppNames)
+	iter := s.StmtStatsIterator(options)
 
-	for _, appName := range appNames {
-		statsContainer := s.getStatsForApplication(appName)
-
-		err := statsContainer.IterateStatementStats(ctx, appName, options.SortedKey, visitor)
-		if err != nil {
-			return fmt.Errorf("sql stats iteration abort: %s", err)
+	for iter.Next() {
+		if err := visitor(ctx, iter.Cur()); err != nil {
+			return err
 		}
 	}
+
 	return nil
 }
 
+// StmtStatsIterator returns an instance of sslocal.StmtStatsIterator for
+// the current SQLStats.
+func (s *SQLStats) StmtStatsIterator(options *sqlstats.IteratorOptions) *StmtStatsIterator {
+	return NewStmtStatsIterator(s, options)
+}
+
 // IterateTransactionStats implements sqlstats.Provider interface.
 func (s *SQLStats) IterateTransactionStats(
 	ctx context.Context, options *sqlstats.IteratorOptions, visitor sqlstats.TransactionVisitor,
 ) error {
-	appNames := s.getAppNames(options.SortedAppNames)
+	iter := s.TxnStatsIterator(options)
 
-	for _, appName := range appNames {
-		statsContainer := s.getStatsForApplication(appName)
-
-		err := statsContainer.IterateTransactionStats(ctx, appName, options.SortedKey, visitor)
-		if err != nil {
-			return fmt.Errorf("sql stats iteration abort: %s", err)
+	for iter.Next() {
+		fingerprintID, stats := iter.Cur()
+		if err := visitor(ctx, fingerprintID, stats); err != nil {
+			return err
 		}
 	}
+
 	return nil
 }
 
+// TxnStatsIterator returns an instance of sslocal.TxnStatsIterator for
+// the current SQLStats.
+func (s *SQLStats) TxnStatsIterator(options *sqlstats.IteratorOptions) *TxnStatsIterator {
+	return NewTxnStatsIterator(s, options)
+}
+
 // IterateAggregatedTransactionStats implements sqlstats.Provider interface.
 func (s *SQLStats) IterateAggregatedTransactionStats(
 	_ context.Context,

pkg/sql/sqlstats/ssmemstorage/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 go_library(
     name = "ssmemstorage",
     srcs = [
+        "ss_mem_iterator.go",
         "ss_mem_storage.go",
         "ss_mem_writer.go",
         "utils.go",

pkg/sql/sqlstats/ssmemstorage/ss_mem_iterator.go

@@ -0,0 +1,169 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package ssmemstorage
+
+import (
+	"sort"
+
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
+)
+
+type baseIterator struct {
+	container *Container
+	idx       int
+}
+
+// StmtStatsIterator is an iterator that iterates over the statement statistics
+// inside of a ssmemstorage.Container.
+type StmtStatsIterator struct {
+	baseIterator
+	stmtKeys stmtList
+}
+
+// NewStmtStatsIterator returns a StmtStatsIterator.
+func NewStmtStatsIterator(
+	container *Container, options *sqlstats.IteratorOptions,
+) *StmtStatsIterator {
+	var stmtKeys stmtList
+	container.mu.Lock()
+	for k := range container.mu.stmts {
+		stmtKeys = append(stmtKeys, k)
+	}
+	container.mu.Unlock()
+	if options.SortedKey {
+		sort.Sort(stmtKeys)
+	}
+
+	return &StmtStatsIterator{
+		baseIterator: baseIterator{
+			container: container,
+			idx:       -1,
+		},
+		stmtKeys: stmtKeys,
+	}
+}
+
+// Next increments the internal counter of the StmtStatsIterator. It returns
+// true if the following Cur() call will be valid, false otherwise.
+func (s *StmtStatsIterator) Next() bool {
+	s.idx++
+	return s.idx < len(s.stmtKeys)
+}
+
+// Cur returns the roachpb.CollectedStatementStatistics at the current internal
+// counter.
+func (s *StmtStatsIterator) Cur() *roachpb.CollectedStatementStatistics {
+	stmtKey := s.stmtKeys[s.idx]
+
+	stmtFingerprintID := constructStatementFingerprintIDFromStmtKey(stmtKey)
+	statementStats, _, _ :=
+		s.container.getStatsForStmtWithKey(stmtKey, invalidStmtFingerprintID, false /* createIfNonexistent */)
+
+	// If the key is not found (and we expected to find it), the table must
+	// have been cleared between now and the time we read all the keys. In
+	// that case we simply skip this key as there are no metrics to report.
+	if statementStats == nil {
+		return nil
+	}
+
+	statementStats.mu.Lock()
+	data := statementStats.mu.data
+	distSQLUsed := statementStats.mu.distSQLUsed
+	vectorized := statementStats.mu.vectorized
+	fullScan := statementStats.mu.fullScan
+	database := statementStats.mu.database
+	statementStats.mu.Unlock()
+
+	collectedStats := roachpb.CollectedStatementStatistics{
+		Key: roachpb.StatementStatisticsKey{
+			Query:       stmtKey.anonymizedStmt,
+			DistSQL:     distSQLUsed,
+			Opt:         true,
+			Vec:         vectorized,
+			ImplicitTxn: stmtKey.implicitTxn,
+			FullScan:    fullScan,
+			Failed:      stmtKey.failed,
+			App:         s.container.appName,
+			Database:    database,
+		},
+		ID:    stmtFingerprintID,
+		Stats: data,
+	}
+
+	return &collectedStats
+}
+
+// TxnStatsIterator is an iterator that iterates over the transaction statistics
+// inside of a ssmemstorage.Container.
+type TxnStatsIterator struct {
+	baseIterator
+	txnKeys txnList
+}
+
+// NewTxnStatsIterator returns a new instance of TxnStatsIterator.
+func NewTxnStatsIterator(
+	container *Container, options *sqlstats.IteratorOptions,
+) *TxnStatsIterator {
+	var txnKeys txnList
+	container.mu.Lock()
+	for k := range container.mu.txns {
+		txnKeys = append(txnKeys, k)
+	}
+	container.mu.Unlock()
+	if options.SortedKey {
+		sort.Sort(txnKeys)
+	}
+
+	return &TxnStatsIterator{
+		baseIterator: baseIterator{
+			container: container,
+			idx:       -1,
+		},
+		txnKeys: txnKeys,
+	}
+}
+
+// Next increments the internal counter of the TxnStatsIterator. It returns
+// true if the following Cur() call will be valid, false otherwise.
+func (t *TxnStatsIterator) Next() bool {
+	t.idx++
+	return t.idx < len(t.txnKeys)
+}
+
+// Cur returns the roachpb.CollectedTransactionStatistics at the current internal
+// counter.
+func (t *TxnStatsIterator) Cur() (
+	roachpb.TransactionFingerprintID,
+	*roachpb.CollectedTransactionStatistics,
+) {
+	txnKey := t.txnKeys[t.idx]
+
+	// We don't want to create the key if it doesn't exist, so it's okay to
+	// pass nil for the statementFingerprintIDs, as they are only set when a key is
+	// constructed.
+	txnStats, _, _ := t.container.getStatsForTxnWithKey(txnKey, nil /* stmtFingerprintIDs */, false /* createIfNonexistent */)
+	// If the key is not found (and we expected to find it), the table must
+	// have been cleared between now and the time we read all the keys. In
+	// that case we simply skip this key as there are no metrics to report.
+	if txnStats == nil {
+		return 0, nil
+	}
+
+	txnStats.mu.Lock()
+	defer txnStats.mu.Unlock()
+	collectedStats := roachpb.CollectedTransactionStatistics{
+		StatementFingerprintIDs: txnStats.statementFingerprintIDs,
+		App:                     t.container.appName,
+		Stats:                   txnStats.mu.data,
+	}
+	return txnKey, &collectedStats
+}