324 lines
10 KiB
Go
324 lines
10 KiB
Go
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])
|
|
}
|