valkey tester improvements

2026-06-09 14:52:14 +02:00
parent 954c86123f
commit ec88214a2e
5 changed files with 426 additions and 0 deletions
@@ -0,0 +1,323 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"fmt"
+	"math/rand/v2"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// keyCount is the number of seed:<n> keys for a given target/value size.
+// Both seed and verify derive it the same way so they agree on the key set
+// without storing a manifest (which the "step 2" actions could mutate).
+func keyCount(targetBytes, valueBytes int) int64 {
+	if valueBytes <= 0 {
+		return 0
+	}
+	return int64((targetBytes + valueBytes - 1) / valueBytes)
+}
+
+// genValue deterministically fills buf with incompressible bytes derived from
+// (runSeed, idx). It's the single source of truth shared by seed and verify:
+// the same (runSeed, idx, len(buf)) always yields identical bytes, which is
+// what lets verify regenerate the expected value instead of remembering it.
+// ChaCha8 gives a fast, high-quality stream so the payload doesn't compress
+// or dedup and thus reflects real memory.
+func genValue(runSeed uint64, idx int64, buf []byte) {
+	var seed [32]byte
+	binary.LittleEndian.PutUint64(seed[0:8], runSeed)
+	binary.LittleEndian.PutUint64(seed[8:16], uint64(idx))
+	rand.NewChaCha8(seed).Read(buf)
+}
+
+// runSeed bulk-loads a single Valkey with deterministic random blobs until the
+// on-wire payload reaches targetBytes. Meant for local use over the zcli VPN
+// (e.g. --host valkey.zerops), so connection params come from the flags.
+//
+// Keys are seed:<n> (n in [0, keyCount)) each holding valueBytes generated by
+// genValue(runSeed, n). Writes are pipelined in batches of batch keys to keep
+// the round-trip count — and VPN latency — from dominating.
+//
+// throttleBytesPerSec, when > 0, caps the average write rate: after each batch
+// the loop sleeps until cumulative bytes / rate seconds have elapsed, so the
+// on-wire throughput converges to the cap without bursting a token bucket.
+func runSeed(opts *redis.Options, runSeed uint64, targetBytes, valueBytes, batch int, throttleBytesPerSec float64) int {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	rdb := redis.NewClient(opts)
+	defer rdb.Close()
+	if err := rdb.Ping(ctx).Err(); err != nil {
+		fmt.Fprintf(os.Stderr, "seed: connect %s failed: %v\n", opts.Addr, err)
+		return 1
+	}
+
+	total := keyCount(targetBytes, valueBytes)
+	throttle := "unlimited"
+	if throttleBytesPerSec > 0 {
+		throttle = humanBytes(int(throttleBytesPerSec)) + "/s"
+	}
+	fmt.Printf("seed: addr=%s db=%d tls=%t target=%s value=%s keys=%d seed=%d batch=%d throttle=%s\n",
+		opts.Addr, opts.DB, opts.TLSConfig != nil,
+		humanBytes(targetBytes), humanBytes(valueBytes), total, runSeed, batch, throttle)
+
+	start, lastLog := time.Now(), time.Now()
+	var written, keys int64
+
+	for keys < total {
+		select {
+		case <-ctx.Done():
+			fmt.Printf("\nseed: interrupted — wrote %d/%d keys (%s)\n", keys, total, time.Since(start).Round(time.Second))
+			return 1
+		default:
+		}
+
+		pipe := rdb.Pipeline()
+		for i := 0; i < batch && keys < total; i++ {
+			// Fresh buffer per key: go-redis keeps the []byte by reference and
+			// only serializes it at Exec, so a reused buffer would store the
+			// batch's last value under every key.
+			buf := make([]byte, valueBytes)
+			genValue(runSeed, keys, buf)
+			pipe.Set(ctx, fmt.Sprintf("seed:%d", keys), buf, 0)
+			keys++
+			written += int64(valueBytes)
+		}
+		if _, err := pipe.Exec(ctx); err != nil {
+			fmt.Fprintf(os.Stderr, "seed: pipeline exec: %v\n", err)
+			return 1
+		}
+
+		// Throttle on cumulative progress: sleep until enough wall-clock has
+		// passed for everything written so far at the target rate. Interruptible
+		// so Ctrl+C doesn't wait out a long sleep.
+		if throttleBytesPerSec > 0 {
+			targetElapsed := time.Duration(float64(written) / throttleBytesPerSec * float64(time.Second))
+			if sleep := targetElapsed - time.Since(start); sleep > 0 {
+				select {
+				case <-ctx.Done():
+				case <-time.After(sleep):
+				}
+			}
+		}
+
+		if now := time.Now(); now.Sub(lastLog) >= time.Second {
+			rate := float64(written) / now.Sub(start).Seconds()
+			fmt.Printf("[%s] %s / %s  (%d/%d keys, %s/s)\n",
+				now.Format("15:04:05"), humanBytes(int(written)), humanBytes(targetBytes), keys, total, humanBytes(int(rate)))
+			lastLog = now
+		}
+	}
+
+	elapsed := time.Since(start)
+	fmt.Printf("seed: done — wrote %s across %d keys in %s (%s/s). verify with the same --seed-mb/--seed-value-bytes/--seed-seed.\n",
+		humanBytes(int(written)), keys, elapsed.Round(time.Millisecond), humanBytes(int(float64(written)/elapsed.Seconds())))
+	return 0
+}
+
+// runVerify regenerates every seed:<n> value and compares it byte-for-byte to
+// what's in Valkey, reporting missing/mismatched keys and any unexpected
+// extras (via DBSIZE). It must be run with the same runSeed/targetBytes/
+// valueBytes used to seed. Returns non-zero if the data is not 1:1.
+func runVerify(opts *redis.Options, runSeed uint64, targetBytes, valueBytes, batch int) int {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	rdb := redis.NewClient(opts)
+	defer rdb.Close()
+	if err := rdb.Ping(ctx).Err(); err != nil {
+		fmt.Fprintf(os.Stderr, "verify: connect %s failed: %v\n", opts.Addr, err)
+		return 1
+	}
+
+	total := keyCount(targetBytes, valueBytes)
+	fmt.Printf("verify: addr=%s db=%d tls=%t expect keys=%d value=%s seed=%d\n",
+		opts.Addr, opts.DB, opts.TLSConfig != nil, total, humanBytes(valueBytes), runSeed)
+
+	start, lastLog := time.Now(), time.Now()
+	var checked, missing, mismatch int64
+	expect := make([]byte, valueBytes)
+
+	for base := int64(0); base < total; base += int64(batch) {
+		select {
+		case <-ctx.Done():
+			fmt.Printf("\nverify: interrupted at %d/%d keys\n", checked, total)
+			return 1
+		default:
+		}
+
+		end := base + int64(batch)
+		if end > total {
+			end = total
+		}
+		pipe := rdb.Pipeline()
+		cmds := make([]*redis.StringCmd, 0, end-base)
+		for idx := base; idx < end; idx++ {
+			cmds = append(cmds, pipe.Get(ctx, fmt.Sprintf("seed:%d", idx)))
+		}
+		if _, err := pipe.Exec(ctx); err != nil && err != redis.Nil {
+			fmt.Fprintf(os.Stderr, "verify: pipeline exec: %v\n", err)
+			return 1
+		}
+
+		for i, cmd := range cmds {
+			idx := base + int64(i)
+			got, err := cmd.Bytes()
+			if err == redis.Nil {
+				missing++
+				if missing <= 10 {
+					fmt.Printf("  MISSING seed:%d\n", idx)
+				}
+				continue
+			}
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "verify: get seed:%d: %v\n", idx, err)
+				return 1
+			}
+			genValue(runSeed, idx, expect)
+			if !bytes.Equal(got, expect) {
+				mismatch++
+				if mismatch <= 10 {
+					fmt.Printf("  MISMATCH seed:%d (len got=%d want=%d)\n", idx, len(got), len(expect))
+				}
+			}
+			checked++
+		}
+
+		if now := time.Now(); now.Sub(lastLog) >= time.Second {
+			fmt.Printf("[%s] verified %d/%d keys\n", now.Format("15:04:05"), base+int64(len(cmds)), total)
+			lastLog = now
+		}
+	}
+
+	dbsize, err := rdb.DBSize(ctx).Result()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "verify: dbsize: %v\n", err)
+		return 1
+	}
+	extras := dbsize - total
+
+	fmt.Printf("verify: %d ok, %d missing, %d mismatched of %d expected; dbsize=%d (%+d vs expected) in %s\n",
+		checked-mismatch, missing, mismatch, total, dbsize, extras, time.Since(start).Round(time.Millisecond))
+
+	if missing == 0 && mismatch == 0 && extras == 0 {
+		fmt.Println("verify: PASS — data is 1:1 with the seed.")
+		return 0
+	}
+	fmt.Println("verify: FAIL — data diverged from the seed.")
+	return 1
+}
+
+// runFlush clears the whole database selected by opts.DB with FLUSHDB,
+// reporting how many keys it dropped. Only the one logical DB is touched —
+// other databases on the same server are left intact.
+func runFlush(opts *redis.Options) int {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	rdb := redis.NewClient(opts)
+	defer rdb.Close()
+	if err := rdb.Ping(ctx).Err(); err != nil {
+		fmt.Fprintf(os.Stderr, "flush: connect %s failed: %v\n", opts.Addr, err)
+		return 1
+	}
+
+	before, err := rdb.DBSize(ctx).Result()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "flush: dbsize db%d: %v\n", opts.DB, err)
+		return 1
+	}
+	if err := rdb.FlushDB(ctx).Err(); err != nil {
+		fmt.Fprintf(os.Stderr, "flush: flushdb db%d: %v\n", opts.DB, err)
+		return 1
+	}
+	fmt.Printf("flush: FLUSHDB db%d on %s — cleared %d keys\n", opts.DB, opts.Addr, before)
+	return 0
+}
+
+// runInspect prints every logical database and its key count. The number of
+// databases comes from CONFIG GET databases (default 16); per-DB key counts
+// come from a single INFO keyspace call, which lists only the non-empty DBs —
+// the rest are reported as 0. Empty DBs are still shown so the full layout is
+// visible at a glance.
+func runInspect(opts *redis.Options) int {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	rdb := redis.NewClient(opts)
+	defer rdb.Close()
+	if err := rdb.Ping(ctx).Err(); err != nil {
+		fmt.Fprintf(os.Stderr, "inspect: connect %s failed: %v\n", opts.Addr, err)
+		return 1
+	}
+
+	numDBs := 16
+	if res, err := rdb.ConfigGet(ctx, "databases").Result(); err == nil {
+		if v, ok := res["databases"]; ok {
+			if n, err := strconv.Atoi(v); err == nil && n > 0 {
+				numDBs = n
+			}
+		}
+	}
+
+	info, err := rdb.Info(ctx, "keyspace").Result()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "inspect: info keyspace: %v\n", err)
+		return 1
+	}
+	// Lines look like: db0:keys=12,expires=0,avg_ttl=0
+	counts := make(map[int]int64)
+	for _, line := range strings.Split(info, "\n") {
+		line = strings.TrimSpace(line)
+		if !strings.HasPrefix(line, "db") {
+			continue
+		}
+		colon := strings.IndexByte(line, ':')
+		if colon < 0 {
+			continue
+		}
+		dbn, err := strconv.Atoi(line[2:colon])
+		if err != nil {
+			continue
+		}
+		for _, field := range strings.Split(line[colon+1:], ",") {
+			if rest, ok := strings.CutPrefix(field, "keys="); ok {
+				if k, err := strconv.ParseInt(rest, 10, 64); err == nil {
+					counts[dbn] = k
+				}
+			}
+		}
+	}
+
+	fmt.Printf("inspect: addr=%s tls=%t databases=%d\n", opts.Addr, opts.TLSConfig != nil, numDBs)
+	var total int64
+	for db := 0; db < numDBs; db++ {
+		k := counts[db]
+		total += k
+		fmt.Printf("  db%-2d  %d keys\n", db, k)
+	}
+	fmt.Printf("inspect: %d keys across %d databases\n", total, numDBs)
+	return 0
+}
+
+func humanBytes(n int) string {
+	const unit = 1024
+	if n < unit {
+		return fmt.Sprintf("%dB", n)
+	}
+	div, exp := int64(unit), 0
+	for v := n / unit; v >= unit; v /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f%cB", float64(n)/float64(div), "KMGTPE"[exp])
+}