roachtest: wait for disk space reclamation in clearrange

In the clearrange roachtests, adapt the completion condition to wait for
reclamation of disk space by waiting until the approximate disk size of the key
range for the dropped table falls to <1% of its original. Additionally,
slightly relax the range count condition to reduce the test time. The long tail
of reducing from a dozen ranges to just 1 can take significant wall time and is
not important to exericse.

Epic: none
Release note: none

Author: jbowens

pkg/cmd/roachtest/tests/clearrange.go

@@ -15,8 +15,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
+	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/stretchr/testify/require"
 )
 
 func registerClearRange(r registry.Registry) {
@@ -82,32 +85,18 @@ func runClearRange(ctx context.Context, t test.Test, c cluster.Cluster, aggressi
 
 	t.Status()
 
-	// Set up a convenience function that we can call to learn the number of
-	// ranges for the bigbank.bank table (even after it's been dropped).
-	numBankRanges := func() func() int {
-		conn := c.Conn(ctx, t.L(), 1)
-		defer conn.Close()
-
-		var startHex string
-		if err := conn.QueryRow(
-			`SELECT to_hex(raw_start_key)
-FROM [SHOW RANGES FROM TABLE bigbank.bank WITH KEYS]
-ORDER BY raw_start_key ASC LIMIT 1`,
-		).Scan(&startHex); err != nil {
-			t.Fatal(err)
-		}
-		return func() int {
-			conn := c.Conn(ctx, t.L(), 1)
-			defer conn.Close()
-			var n int
-			if err := conn.QueryRow(
-				`SELECT count(*) FROM crdb_internal.ranges_no_leases WHERE substr(to_hex(start_key), 1, length($1::string)) = $1`, startHex,
-			).Scan(&n); err != nil {
-				t.Fatal(err)
-			}
-			return n
-		}
-	}()
+	bigBankSpan, err := getKeyspanForTable(ctx, t, c, 1, "bigbank.bank")
+	require.NoError(t, err)
+	t.L().Printf("bigbank DB ID: %s (%x - %x)", bigBankSpan, bigBankSpan.Key, bigBankSpan.EndKey)
+	getBigBankStats := func() spanStats {
+		stats := getSpanStats(ctx, t, c, 1, bigBankSpan)
+		t.L().Printf("bigbank: %d ranges, %s disk, %s live, %s total",
+			stats.rangeCount,
+			humanizeutil.IBytes(stats.approximateDiskBytes),
+			humanizeutil.IBytes(stats.liveBytes),
+			humanizeutil.IBytes(stats.totalBytes))
+		return stats
+	}
 
 	m.Go(func(ctx context.Context) error {
 		c.Run(ctx, option.WithNodes(c.Node(1)), `./cockroach workload init kv {pgurl:1}`)
@@ -127,6 +116,10 @@ ORDER BY raw_start_key ASC LIMIT 1`,
 			return err
 		}
 
+		// Collect the stats before dropping the table. getBigBanksStats will
+		// print them out.
+		_ = getBigBankStats()
+
 		t.WorkerStatus("dropping table")
 		defer t.WorkerStatus()
 
@@ -136,26 +129,39 @@ ORDER BY raw_start_key ASC LIMIT 1`,
 			return err
 		}
 
-		t.WorkerStatus("computing number of ranges")
-		initialBankRanges := numBankRanges()
+		t.WorkerStatus("computing span stats")
+		preDropBankStats := getBigBankStats()
 
 		t.WorkerStatus("dropping bank table")
 		if _, err := conn.ExecContext(ctx, `DROP TABLE bigbank.bank`); err != nil {
 			return err
 		}
 
-		// Spend some time reading data with a timeout to make sure the
-		// DROP above didn't brick the cluster. At the time of writing,
-		// clearing all of the table data takes ~6min, so we want to run
-		// for at least a multiple of that duration.
-		const minDuration = 45 * time.Minute
-		deadline := timeutil.Now().Add(minDuration)
-		curBankRanges := numBankRanges()
-		t.WorkerStatus("waiting for ~", curBankRanges, " merges to complete (and for at least ", minDuration, " to pass)")
-		for timeutil.Now().Before(deadline) || curBankRanges > 1 {
-			after := time.After(5 * time.Minute)
-			curBankRanges = numBankRanges() // this call takes minutes, unfortunately
-			t.WorkerProgress(1 - float64(curBankRanges)/float64(initialBankRanges))
+		curBankStats := getBigBankStats()
+		progressFn := func() float64 {
+			// Compute progress as a float [0, 1.0].
+			//
+			// We compute the progress in terms of the number of ranges and the
+			// amount of disk space, relative to the stats we computed
+			// immediately after dropping the table. That is:
+			//
+			//    1 - (current / initial)
+			//
+			// The range count progress subtracts 1 from each count to account
+			// for the expectation that one range will always remain.
+			//
+			// We compute the overall progress as the minimum of the two metrics'
+			// progress.
+			mergeProgress := 1 - (float64(curBankStats.rangeCount-1) /
+				float64(preDropBankStats.rangeCount-1))
+			diskProgress := 1 - (float64(curBankStats.approximateDiskBytes) /
+				float64(preDropBankStats.approximateDiskBytes))
+			return max(0, min(mergeProgress, diskProgress))
+		}
+		// Terminate when progress is 0.975 or greater. That is, we've reclaimed
+		// 97.5% of the disk space and merged 97.5% of the ranges.
+		for progress := progressFn(); progress < 0.975; progress = progressFn() {
+			t.WorkerProgress(progress)
 
 			var count int
 			// NB: context cancellation in QueryRowContext does not work as expected.
@@ -168,16 +174,65 @@ ORDER BY raw_start_key ASC LIMIT 1`,
 				return err
 			}
 
-			t.WorkerStatus("waiting for ~", curBankRanges, " merges to complete (and for at least ", timeutil.Until(deadline), " to pass)")
+			t.WorkerStatus("progress ", progress, " (", curBankStats.rangeCount, " ranges, ",
+				humanizeutil.IBytes(curBankStats.approximateDiskBytes), " disk usage)")
 			select {
-			case <-after:
+			case <-time.After(time.Minute):
 			case <-ctx.Done():
 				return ctx.Err()
 			}
+			curBankStats = getBigBankStats()
 		}
-		// TODO(tschottdorf): verify that disk space usage drops below to <some small amount>, but that
-		// may not actually happen (see https://github.com/cockroachdb/cockroach/issues/29290).
+		t.WorkerStatus("reclamation condition met")
 		return nil
 	})
 	m.Wait()
 }
+
+func getKeyspanForTable(
+	ctx context.Context, t test.Test, c cluster.Cluster, n int, tbl string,
+) (roachpb.Span, error) {
+	conn := c.Conn(ctx, t.L(), n)
+	defer conn.Close()
+	var startKey, endKey roachpb.Key
+	err := conn.QueryRow(`SELECT `+
+		`crdb_internal.table_span($1::regclass::oid::int)[1] AS start_key, `+
+		`crdb_internal.table_span($1::regclass::oid::int)[2] AS end_key`,
+		tbl).Scan(&startKey, &endKey)
+	return roachpb.Span{
+		Key:    startKey,
+		EndKey: endKey,
+	}, err
+}
+
+type spanStats struct {
+	rangeCount           int
+	approximateDiskBytes int64
+	liveBytes            int64
+	totalBytes           int64
+}
+
+func getSpanStats(
+	ctx context.Context, t test.Test, c cluster.Cluster, n int, span roachpb.Span,
+) spanStats {
+	conn := c.Conn(ctx, t.L(), n)
+	defer conn.Close()
+
+	var stats spanStats
+	err := conn.QueryRow(
+		`SELECT `+
+			`(stats->'range_count')::int AS range_count, `+
+			`(stats->'approximate_disk_bytes')::int AS approximate_disk_bytes, `+
+			`(stats->'approximate_total_stats'->'live_bytes')::int AS live_bytes, `+
+			`(stats->'approximate_total_stats'->'key_bytes')::int + (stats->'approximate_total_stats'->'val_bytes')::int AS total_bytes `+
+			`FROM crdb_internal.tenant_span_stats(ARRAY(SELECT($1::bytes, $2::bytes)))`,
+		span.Key, span.EndKey).
+		Scan(
+			&stats.rangeCount,
+			&stats.approximateDiskBytes,
+			&stats.liveBytes,
+			&stats.totalBytes,
+		)
+	require.NoError(t, err)
+	return stats
+}