Merge #107120 #107248

107120: server,storage: breakdown byte counts in SpanStats by backing storage r=dt a=dt

Release note: none.
Epic: none.

107248: rpc: deflake `TestInitialHeartbeatFailedError` r=erikgrinaker a=erikgrinaker

Resolves #107245.
Epic: none
Release note: None

Co-authored-by: David Taylor <[email protected]>
Co-authored-by: Erik Grinaker <[email protected]>

pkg/cli/debug.go

@@ -936,7 +936,7 @@ func runDebugCompact(cmd *cobra.Command, args []string) error {
 	}
 
 	{
-		approxBytesBefore, err := db.ApproximateDiskBytes(roachpb.KeyMin, roachpb.KeyMax)
+		approxBytesBefore, _, _, err := db.ApproximateDiskBytes(roachpb.KeyMin, roachpb.KeyMax)
 		if err != nil {
 			return errors.Wrap(err, "while computing approximate size before compaction")
 		}
@@ -966,7 +966,7 @@ func runDebugCompact(cmd *cobra.Command, args []string) error {
 	fmt.Printf("%s\n", db.GetMetrics())
 
 	{
-		approxBytesAfter, err := db.ApproximateDiskBytes(roachpb.KeyMin, roachpb.KeyMax)
+		approxBytesAfter, _, _, err := db.ApproximateDiskBytes(roachpb.KeyMin, roachpb.KeyMax)
 		if err != nil {
 			return errors.Wrap(err, "while computing approximate size after compaction")
 		}

pkg/kv/kvserver/replica.go

@@ -2336,7 +2336,8 @@ func (r *Replica) GetEngineCapacity() (roachpb.StoreCapacity, error) {
 // GetApproximateDiskBytes returns an approximate measure of bytes in the store
 // in the specified key range.
 func (r *Replica) GetApproximateDiskBytes(from, to roachpb.Key) (uint64, error) {
-	return r.store.TODOEngine().ApproximateDiskBytes(from, to)
+	bytes, _, _, err := r.store.TODOEngine().ApproximateDiskBytes(from, to)
+	return bytes, err
 }
 
 func init() {

pkg/roachpb/span_stats.go

@@ -56,3 +56,10 @@ var RangeDescPageSize = settings.RegisterIntSetting(
 		return nil
 	},
 )
+
+func (m *SpanStats) Add(other *SpanStats) {
+	m.TotalStats.Add(other.TotalStats)
+	m.ApproximateDiskBytes += other.ApproximateDiskBytes
+	m.RemoteFileBytes += other.RemoteFileBytes
+	m.ExternalFileBytes += other.ExternalFileBytes
+}

pkg/roachpb/span_stats.proto

@@ -43,14 +43,29 @@ message SpanStats {
   // range end, will have a range_count value of 1.
   int32 range_count = 2;
 
-  // The explicit jsontag prevents 'omitempty` from being added by default.
+  // ApproximateDiskBytes is the approximate size "on-disk" in all files of the
+  // data in the span. NB; this *includes* files stored remotely, not just on
+  // _local_ disk; see the RemoteFileBytes field below.
+  // NB: The explicit jsontag prevents 'omitempty` from being added by default.
   uint64 approximate_disk_bytes = 3 [(gogoproto.jsontag) = "approximate_disk_bytes"];
+
+  // RemoteFileBytes is the subset of ApproximateDiskBytes which are stored in
+  // "remote" files (i.e. shared files and external files).
+  uint64 remote_file_bytes = 5;
+
+  // ExternalFileBytes is the subset of RemoteFileBytes that are in "external"
+  // files (not written/owned by this cluster, such as in restored backups).
+  uint64 external_file_bytes = 6;
+
+  // NEXT ID: 7.
 }
 
 message SpanStatsResponse {
-  cockroach.storage.enginepb.MVCCStats total_stats = 1 [(gogoproto.nullable) = false];
-  // See the range_count comment for the SpanStats proto.
-  int32 range_count = 2;
-  uint64 approximate_disk_bytes = 3;
+  reserved 1;
+  reserved 2;
+  reserved 3;
+
   map<string, SpanStats> span_to_stats = 4;
+
+  // NEXT ID: 5.
 }

pkg/rpc/context_test.go

@@ -2556,7 +2556,9 @@ func TestInitialHeartbeatFailedError(t *testing.T) {
 	_, err = clientCtx.GRPCDialNode(remoteAddr, nodeID, SystemClass).Connect(ctx)
 	requireHeartbeatError(t, err)
 
-	// Start server listener.
+	// Start server listener. We set the ping handler to fail initially, to make
+	// sure no other actor creates an RPC connection.
+	failPing.Store(true)
 	s := newTestServer(t, serverCtx)
 	RegisterHeartbeatServer(s, &HeartbeatService{
 		clock:              clock,
@@ -2582,10 +2584,12 @@ func TestInitialHeartbeatFailedError(t *testing.T) {
 	failPing.Store(true)
 	_, err = clientCtx.GRPCDialNode(remoteAddr, nodeID, SystemClass).Connect(ctx)
 	requireHeartbeatError(t, err)
-	failPing.Store(false)
 
-	// Stalled pings result in InitialHeartbeatFailedError.
+	// Stalled pings result in InitialHeartbeatFailedError. We're careful to
+	// enable the hang before disabling the error, to make sure no other
+	// actor establishes a connection in the meanwhile.
 	hangPing.Store(true)
+	failPing.Store(false)
 	_, err = clientCtx.GRPCDialNode(remoteAddr, nodeID, SystemClass).Connect(ctx)
 	requireHeartbeatError(t, err)
 	hangPing.Store(false)

pkg/server/span_stats_server.go

@@ -85,8 +85,7 @@ func (s *systemStatusServer) spanStatsFanOut(
 			if !exists {
 				res.SpanToStats[spanStr] = spanStats
 			} else {
-				res.SpanToStats[spanStr].ApproximateDiskBytes += spanStats.ApproximateDiskBytes
-				res.SpanToStats[spanStr].TotalStats.Add(spanStats.TotalStats)
+				res.SpanToStats[spanStr].Add(spanStats)
 			}
 		}
 	}
@@ -218,11 +217,13 @@ func (s *systemStatusServer) statsForSpan(
 	}
 	// Finally, get the approximate disk bytes from each store.
 	err = s.stores.VisitStores(func(store *kvserver.Store) error {
-		approxDiskBytes, err := store.TODOEngine().ApproximateDiskBytes(rSpan.Key.AsRawKey(), rSpan.EndKey.AsRawKey())
+		approxDiskBytes, remoteBytes, externalBytes, err := store.TODOEngine().ApproximateDiskBytes(rSpan.Key.AsRawKey(), rSpan.EndKey.AsRawKey())
 		if err != nil {
 			return err
 		}
 		spanStats.ApproximateDiskBytes += approxDiskBytes
+		spanStats.RemoteFileBytes += remoteBytes
+		spanStats.ExternalFileBytes += externalBytes
 		return nil
 	})
 

pkg/storage/engine.go

@@ -988,8 +988,11 @@ type Engine interface {
 	// When called, it may choose to block if the engine determines that it is in
 	// or approaching a state where further ingestions may risk its health.
 	PreIngestDelay(ctx context.Context)
-	// ApproximateDiskBytes returns an approximation of the on-disk size for the given key span.
-	ApproximateDiskBytes(from, to roachpb.Key) (uint64, error)
+	// ApproximateDiskBytes returns an approximation of the on-disk size and file
+	// counts for the given key span, along with how many of those bytes are on
+	// remote, as well as specifically external remote, storage.
+	ApproximateDiskBytes(from, to roachpb.Key) (total, remote, external uint64, _ error)
+
 	// CompactRange ensures that the specified range of key value pairs is
 	// optimized for space efficiency.
 	CompactRange(start, end roachpb.Key) error

pkg/storage/pebble.go

@@ -2040,14 +2040,16 @@ func (p *Pebble) GetTableMetrics(start, end roachpb.Key) ([]enginepb.SSTableMetr
 }
 
 // ApproximateDiskBytes implements the Engine interface.
-func (p *Pebble) ApproximateDiskBytes(from, to roachpb.Key) (uint64, error) {
+func (p *Pebble) ApproximateDiskBytes(
+	from, to roachpb.Key,
+) (bytes, remoteBytes, externalBytes uint64, _ error) {
 	fromEncoded := EngineKey{Key: from}.Encode()
 	toEncoded := EngineKey{Key: to}.Encode()
-	count, err := p.db.EstimateDiskUsage(fromEncoded, toEncoded)
+	bytes, remoteBytes, externalBytes, err := p.db.EstimateDiskUsageByBackingType(fromEncoded, toEncoded)
 	if err != nil {
-		return 0, err
+		return 0, 0, 0, err
 	}
-	return count, nil
+	return bytes, remoteBytes, externalBytes, nil
 }
 
 // Compact implements the Engine interface.

pkg/storage/pebble_test.go

@@ -1374,7 +1374,7 @@ func TestApproximateDiskBytes(t *testing.T) {
 	require.NoError(t, p.Flush())
 
 	approxBytes := func(span roachpb.Span) uint64 {
-		v, err := p.ApproximateDiskBytes(span.Key, span.EndKey)
+		v, _, _, err := p.ApproximateDiskBytes(span.Key, span.EndKey)
 		require.NoError(t, err)
 		t.Logf("%s (%x-%x): %d bytes", span, span.Key, span.EndKey, v)
 		return v