feat(bench): add corpus-gen progress logging and speed up sampling

bench/rpc/cmd/corpus-gen/class_hash.go

@@ -38,11 +38,11 @@
 	if err != nil {
 		return "", err
 	}
-	class, err := input.client.classAt(input.ctx, blockNumber, classHash)
+	isSierra, err := input.cache.isSierra(input.ctx, input.client, blockNumber, classHash)
 	if err != nil {
 		return "", err
 	}
-	if len(class.SierraProgram) == 0 {
+	if !isSierra {
 		return "", errResample
 	}
 	return classHash, nil

bench/rpc/cmd/corpus-gen/get_storage_proof.go

@@ -46,28 +46,36 @@
 	// sampled trie members persist, so historical diffs are still valid sources.
 	params := storageProofParams{BlockID: "latest"}
 	for range input.args.NumClasses {
-		classHash, err := sampleSierraClassHash(input, input.args.sampleBlockNumber(input.rng))
+		// Cairo 0 hashes are fine here: keys absent from the classes trie
+		// yield valid non-membership proofs with near-identical node work.
+		classHash, err := resample(func() (string, error) {
+			return sampleClassHash(input, input.args.sampleBlockNumber(input.rng))
+		})
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("sample class hash: %w", err)
 		}
 		params.ClassHashes = append(params.ClassHashes, classHash)
 	}
 	for range input.args.NumContracts {
-		address, err := sampleContractAddress(
-			input,
-			input.args.sampleBlockNumber(input.rng),
-			storageDiffAddresses,
-		)
+		address, err := resample(func() (string, error) {
+			return sampleContractAddress(
+				input,
+				input.args.sampleBlockNumber(input.rng),
+				storageDiffAddresses,
+			)
+		})
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("sample contract address: %w", err)
 		}
 		params.ContractAddresses = append(params.ContractAddresses, address)
 	}
 	keyIndex := make(map[string]int)
 	for range input.args.NumKeys {
-		entry, err := sampleStorageEntry(input, input.args.sampleBlockNumber(input.rng))
+		entry, err := resample(func() (storageEntry, error) {
+			return sampleStorageEntry(input, input.args.sampleBlockNumber(input.rng))
+		})
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("sample storage key: %w", err)
 		}
 		if i, ok := keyIndex[entry.Address]; ok {
 			params.ContractsStorageKeys[i].StorageKeys = append(

bench/rpc/cmd/corpus-gen/main.go

@@ -8,6 +8,8 @@
 	"math/rand/v2"
 	"os"
 	"runtime"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/sourcegraph/conc/pool"
@@ -166,7 +168,9 @@
 ) error {
 	generatedAt := time.Now().UTC().Format(time.RFC3339)
 
-	c, err := buildCorpus(cmd.Context(), cfg, client, method, meta, generatedAt, gen)
+	c, err := buildCorpus(
+		cmd.Context(), cfg, client, method, meta, generatedAt, gen, cmd.ErrOrStderr(),
+	)
 	if err != nil {
 		return err
 	}
@@ -187,12 +191,17 @@
 	meta T,
 	generatedAt string,
 	gen paramsGen,
+	progress io.Writer,
 ) (*corpus[T], error) {
 	version, err := client.specVersion(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("fetch spec version: %w", err)
 	}
 
+	var completed atomic.Int64
+	stopProgress := reportProgress(progress, &completed, cfg.count)
+	defer stopProgress()
+
 	requests := make([]any, cfg.count)
 	p := pool.New().
 		WithContext(ctx).
@@ -215,9 +224,10 @@
 			}
 			if cfg.batch == 0 {
 				requests[i] = entry[0]
-				return nil
+			} else {
+				requests[i] = entry
 			}
-			requests[i] = entry
+			completed.Add(1)
 			return nil
 		})
 	}
@@ -239,6 +249,31 @@
 	}, nil
 }
 
+const progressInterval = 2 * time.Second
+
+func reportProgress(w io.Writer, completed *atomic.Int64, total int) (stop func()) {
+	quit := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ticker := time.NewTicker(progressInterval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-quit:
+				return
+			case <-ticker.C:
+				fmt.Fprintf(w, "progress: %d/%d entries\n", completed.Load(), total)
+			}
+		}
+	}()
+	return func() {
+		close(quit)
+		wg.Wait()
+	}
+}
+
 func newSeededRand(seed, stream uint64) *rand.Rand {
 	//nolint:gosec // G404: deterministic corpus needs a seeded PRNG, not crypto randomness
 	return rand.New(rand.NewPCG(seed, stream))

bench/rpc/cmd/corpus-gen/sampler.go

@@ -17,6 +17,7 @@
 	client *rpcClient
 	rng    *rand.Rand
 	args   *T
+	cache  *sierraCache
 }
 
 type sampler[T any] func(input samplerInput[T]) (any, error)
@@ -54,32 +55,44 @@
 				return err
 			}
 		}
+		cache := newSierraCache()
 		gen := func(ctx context.Context, client *rpcClient, rng *rand.Rand) (any, error) {
 			input := samplerInput[T]{
 				ctx:    ctx,
 				client: client,
 				rng:    rng,
 				args:   args,
+				cache:  cache,
 			}
-			for range maxResampleAttempts {
-				result, err := sample(input)
-				if err == nil {
-					return result, nil
-				}
-				if !errors.Is(err, errResample) {
-					return nil, err
-				}
+			result, err := resample(func() (any, error) { return sample(input) })
+			if err != nil {
+				return nil, fmt.Errorf("%s: %w", method, err)
 			}
-			return nil, fmt.Errorf("%s: no candidate found after %d attempts", method, maxResampleAttempts)
+			return result, nil
 		}
 		return runCorpus(cmd, cfg, client, method, args, gen)
 	}
 	return cmd
 }
 
-// errResample tells newSampledCmd's loop to re-invoke the sampler.
+// errResample tells resample to retry fn with fresh draws.
 var errResample = errors.New("resample")
 
+// resample retries fn until it succeeds or fails with a non-errResample error.
+func resample[R any](fn func() (R, error)) (R, error) {
+	var zero R
+	for range maxResampleAttempts {
+		result, err := fn()
+		if err == nil {
+			return result, nil
+		}
+		if !errors.Is(err, errResample) {
+			return zero, err
+		}
+	}
+	return zero, fmt.Errorf("no candidate found after %d attempts", maxResampleAttempts)
+}
+
 func commandName(method string) string {
 	return strings.TrimPrefix(method, "starknet_")
 }

bench/rpc/cmd/corpus-gen/sierra_cache.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"context"
+	"sync"
+
+	"golang.org/x/sync/singleflight"
+)
+
+// sierraCache memoizes per class hash whether the class is Sierra, sparing
+// repeat getClass downloads; a class's content is immutable, so verdicts
+// never expire. Errors are not cached.
+type sierraCache struct {
+	group    singleflight.Group
+	mu       sync.RWMutex
+	verdicts map[string]bool
+}
+
+func newSierraCache() *sierraCache {
+	return &sierraCache{verdicts: make(map[string]bool)}
+}
+
+func (c *sierraCache) isSierra(
+	ctx context.Context,
+	client *rpcClient,
+	blockNumber uint64,
+	classHash string,
+) (bool, error) {
+	if verdict, ok := c.get(classHash); ok {
+		return verdict, nil
+	}
+	result, err, _ := c.group.Do(classHash, func() (any, error) {
+		class, err := client.classAt(ctx, blockNumber, classHash)
+		if err != nil {
+			return false, err
+		}
+		isSierra := len(class.SierraProgram) > 0
+		c.set(classHash, isSierra)
+		return isSierra, nil
+	})
+	if err != nil {
+		return false, err
+	}
+	return result.(bool), nil
+}
+
+func (c *sierraCache) get(classHash string) (verdict, ok bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	verdict, ok = c.verdicts[classHash]
+	return verdict, ok
+}
+
+func (c *sierraCache) set(classHash string, isSierra bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.verdicts[classHash] = isSierra
+}