github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/spanner/stats.go (about)

     1  package spanner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"cloud.google.com/go/spanner"
     9  	"go.opentelemetry.io/otel/trace"
    10  
    11  	"github.com/authzed/spicedb/pkg/datastore"
    12  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
    13  )
    14  
    15  var querySomeRandomRelationships = fmt.Sprintf(`SELECT %s FROM %s LIMIT 10`,
    16  	strings.Join([]string{
    17  		colNamespace,
    18  		colObjectID,
    19  		colRelation,
    20  		colUsersetNamespace,
    21  		colUsersetObjectID,
    22  		colUsersetRelation,
    23  	}, ", "),
    24  	tableRelationship)
    25  
    26  const defaultEstimatedBytesPerRelationships = 20 // determined by looking at some sample clusters
    27  
    28  func (sd *spannerDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
    29  	var uniqueID string
    30  	if err := sd.client.Single().Read(
    31  		context.Background(),
    32  		tableMetadata,
    33  		spanner.AllKeys(),
    34  		[]string{colUniqueID},
    35  	).Do(func(r *spanner.Row) error {
    36  		return r.Columns(&uniqueID)
    37  	}); err != nil {
    38  		return datastore.Stats{}, fmt.Errorf("unable to read unique ID: %w", err)
    39  	}
    40  
    41  	iter := sd.client.Single().Read(
    42  		ctx,
    43  		tableNamespace,
    44  		spanner.AllKeys(),
    45  		[]string{colNamespaceConfig, colNamespaceTS},
    46  	)
    47  	defer iter.Stop()
    48  
    49  	allNamespaces, err := readAllNamespaces(iter, trace.SpanFromContext(ctx))
    50  	if err != nil {
    51  		return datastore.Stats{}, fmt.Errorf("unable to read namespaces: %w", err)
    52  	}
    53  
    54  	// If there is not yet a cached estimated bytes per relationship, read a few relationships and then
    55  	// compute the average bytes per relationship.
    56  	sd.cachedEstimatedBytesPerRelationshipLock.RLock()
    57  	estimatedBytesPerRelationship := sd.cachedEstimatedBytesPerRelationship
    58  	sd.cachedEstimatedBytesPerRelationshipLock.RUnlock()
    59  
    60  	if estimatedBytesPerRelationship == 0 {
    61  		riter := sd.client.Single().Query(ctx, spanner.Statement{SQL: querySomeRandomRelationships})
    62  		defer riter.Stop()
    63  
    64  		totalByteCount := 0
    65  		totalRelationships := 0
    66  
    67  		if err := riter.Do(func(row *spanner.Row) error {
    68  			nextTuple := &core.RelationTuple{
    69  				ResourceAndRelation: &core.ObjectAndRelation{},
    70  				Subject:             &core.ObjectAndRelation{},
    71  			}
    72  			err := row.Columns(
    73  				&nextTuple.ResourceAndRelation.Namespace,
    74  				&nextTuple.ResourceAndRelation.ObjectId,
    75  				&nextTuple.ResourceAndRelation.Relation,
    76  				&nextTuple.Subject.Namespace,
    77  				&nextTuple.Subject.ObjectId,
    78  				&nextTuple.Subject.Relation,
    79  			)
    80  			if err != nil {
    81  				return err
    82  			}
    83  
    84  			relationshipByteCount := len(nextTuple.ResourceAndRelation.Namespace) + len(nextTuple.ResourceAndRelation.ObjectId) +
    85  				len(nextTuple.ResourceAndRelation.Relation) + len(nextTuple.Subject.Namespace) + len(nextTuple.Subject.ObjectId) +
    86  				len(nextTuple.Subject.Relation)
    87  
    88  			totalRelationships++
    89  			totalByteCount += relationshipByteCount
    90  
    91  			return nil
    92  		}); err != nil {
    93  			return datastore.Stats{}, err
    94  		}
    95  
    96  		if totalRelationships == 0 {
    97  			return datastore.Stats{
    98  				UniqueID:                   uniqueID,
    99  				ObjectTypeStatistics:       datastore.ComputeObjectTypeStats(allNamespaces),
   100  				EstimatedRelationshipCount: 0,
   101  			}, nil
   102  		}
   103  
   104  		estimatedBytesPerRelationship = uint64(totalByteCount / totalRelationships)
   105  		if estimatedBytesPerRelationship > 0 {
   106  			sd.cachedEstimatedBytesPerRelationshipLock.Lock()
   107  			sd.cachedEstimatedBytesPerRelationship = estimatedBytesPerRelationship
   108  			sd.cachedEstimatedBytesPerRelationshipLock.Unlock()
   109  		}
   110  	}
   111  
   112  	if estimatedBytesPerRelationship == 0 {
   113  		estimatedBytesPerRelationship = defaultEstimatedBytesPerRelationships // Use a default
   114  	}
   115  
   116  	// Reference: https://cloud.google.com/spanner/docs/introspection/table-sizes-statistics
   117  	queryRelationshipByteEstimate := fmt.Sprintf(`SELECT used_bytes FROM %s WHERE
   118  		interval_end = (
   119  			SELECT MAX(interval_end)
   120  			FROM %s
   121  		)
   122  		AND table_name = '%s'`, sd.tableSizesStatsTable, sd.tableSizesStatsTable, tableRelationship)
   123  
   124  	var byteEstimate spanner.NullInt64
   125  	if err := sd.client.Single().Query(ctx, spanner.Statement{SQL: queryRelationshipByteEstimate}).Do(func(r *spanner.Row) error {
   126  		return r.Columns(&byteEstimate)
   127  	}); err != nil {
   128  		return datastore.Stats{}, fmt.Errorf("unable to read tuples byte count: %w", err)
   129  	}
   130  
   131  	// If the byte estimate is NULL, try to fallback to just selecting the single row. This is necessary for certain
   132  	// versions of the emulator.
   133  	if byteEstimate.IsNull() {
   134  		lookupSingleEstimate := fmt.Sprintf(`SELECT used_bytes FROM %s WHERE table_name = '%s'`, sd.tableSizesStatsTable, tableRelationship)
   135  		if err := sd.client.Single().Query(ctx, spanner.Statement{SQL: lookupSingleEstimate}).Do(func(r *spanner.Row) error {
   136  			return r.Columns(&byteEstimate)
   137  		}); err != nil {
   138  			return datastore.Stats{}, fmt.Errorf("unable to fallback read tuples byte count: %w", err)
   139  		}
   140  	}
   141  
   142  	return datastore.Stats{
   143  		UniqueID:                   uniqueID,
   144  		ObjectTypeStatistics:       datastore.ComputeObjectTypeStats(allNamespaces),
   145  		EstimatedRelationshipCount: uint64(byteEstimate.Int64) / estimatedBytesPerRelationship,
   146  	}, nil
   147  }