roachtest: validate token return in perturbation/* tests

This commit adds validation that all RAC tokens are returned on all
stable nodes at the end of the test.

Fixes: cockroachdb#133410

Release note: None

pkg/cmd/roachtest/tests/admission_control_latency.go

@@ -27,6 +27,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -901,6 +902,47 @@ func (v variations) runTest(ctx context.Context, t test.Test, c cluster.Cluster)
 	t.L().Printf("validating stats after the perturbation")
 	failures = append(failures, isAcceptableChange(t.L(), baselineStats, afterStats, v.acceptableChange)...)
 	require.True(t, len(failures) == 0, strings.Join(failures, "\n"))
+	v.validateTokensReturned(ctx, t)
+}
+
+// validateTokensReturned ensures that all RAC tokens are returned to the pool
+// at the end of the test.
+func (v variations) validateTokensReturned(ctx context.Context, t test.Test) {
+	t.L().Printf("validating all tokens returned")
+	for _, node := range v.stableNodes() {
+		// Wait for the tokens to be returned to the pool. Normally this will
+		// pass immediately however it is possible that there is still some
+		// recovery so loop a few times.
+		testutils.SucceedsWithin(t, func() error {
+			db := v.Conn(ctx, t.L(), node)
+			defer db.Close()
+			for _, sType := range []string{"send", "eval"} {
+				for _, tType := range []string{"elastic", "regular"} {
+					statPrefix := fmt.Sprintf("kvflowcontrol.tokens.%s.%s", sType, tType)
+					query := fmt.Sprintf(`
+		SELECT d.value::INT8 AS deducted, r.value::INT8 AS returned
+		FROM
+		  crdb_internal.node_metrics d,
+		  crdb_internal.node_metrics r
+		WHERE
+		  d.name='%s.deducted' AND
+		  r.name='%s.returned'`,
+						statPrefix, statPrefix)
+					rows, err := db.QueryContext(ctx, query)
+					require.NoError(t, err)
+					require.True(t, rows.Next())
+					var deducted, returned int64
+					if err := rows.Scan(&deducted, &returned); err != nil {
+						return err
+					}
+					if deducted != returned {
+						return errors.Newf("tokens not returned for %s: deducted %d returned %d", statPrefix, deducted, returned)
+					}
+				}
+			}
+			return nil
+		}, 5*time.Second)
+	}
 }
 
 // trackedStat is a collection of the relevant values from the histogram. The