From efc85eccfff41c9f8aabb80d2d727361c7e371b9 Mon Sep 17 00:00:00 2001 From: bdchatham Date: Thu, 11 Jun 2026 15:40:15 -0700 Subject: [PATCH] feat(config): add Distribution primitive + Name discriminator (PLT-455) Add a tagged Distribution config type mirroring GasPicker: an UnmarshalJSON switch on a "Name" discriminator selecting "uniform" or "zipfian" delegates, with a loud default error on unknown kinds. The zipfian "theta" parameter is range-checked at parse time. Wire optional KeyDistribution and SizeDistribution onto Scenario as additive, omitempty *Distribution fields so existing profiles round-trip unchanged. Sampling math (YCSB precomputed-zeta zipfian + seeded uniform draw) is deferred to PLT-460; SampleIndex methods are placeholders. Co-Authored-By: Claude Opus 4.8 (1M context) --- config/config.go | 14 ++--- config/distribution.go | 102 ++++++++++++++++++++++++++++++++++++ config/distribution_test.go | 68 ++++++++++++++++++++++++ 3 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 config/distribution.go create mode 100644 config/distribution_test.go diff --git a/config/config.go b/config/config.go index 74105cf..7862e13 100644 --- a/config/config.go +++ b/config/config.go @@ -61,10 +61,12 @@ type AccountConfig struct { // Scenario represents each scenario in the load configuration. type Scenario struct { - Name string `json:"name,omitempty"` - Weight int `json:"weight,omitempty"` - Accounts *AccountConfig `json:"accounts,omitempty"` - GasPicker *GasPicker `json:"gasPicker,omitempty"` - GasFeeCapPicker *GasPicker `json:"gasFeeCapPicker,omitempty"` - GasTipCapPicker *GasPicker `json:"gasTipCapPicker,omitempty"` + Name string `json:"name,omitempty"` + Weight int `json:"weight,omitempty"` + Accounts *AccountConfig `json:"accounts,omitempty"` + GasPicker *GasPicker `json:"gasPicker,omitempty"` + GasFeeCapPicker *GasPicker `json:"gasFeeCapPicker,omitempty"` + GasTipCapPicker *GasPicker `json:"gasTipCapPicker,omitempty"` + KeyDistribution *Distribution `json:"keyDistribution,omitempty"` + SizeDistribution *Distribution `json:"sizeDistribution,omitempty"` } diff --git a/config/distribution.go b/config/distribution.go new file mode 100644 index 0000000..15585bb --- /dev/null +++ b/config/distribution.go @@ -0,0 +1,102 @@ +package config + +import ( + "encoding/json" + "fmt" +) + +var ( + _ indexSampler = (*Distribution)(nil) + _ indexSampler = (*UniformDistribution)(nil) + _ indexSampler = (*ZipfianDistribution)(nil) +) + +// indexSampler draws an index in [0, n) from some keyspace distribution. +type indexSampler interface { + SampleIndex(n uint64) (uint64, error) +} + +// Distribution is a tagged wrapper over a keyspace index distribution, selected +// by a "Name" discriminator on the JSON wire format. The discriminator strings +// ("uniform", "zipfian") and the "theta" parameter name are a frozen +// saved-workload contract; do not rename them. +type Distribution struct { + name string + delegate indexSampler +} + +func (d *Distribution) Name() string { return d.name } + +// SampleIndex delegates to the selected distribution. A zero-value (no Name) +// Distribution samples nothing and returns 0. +func (d *Distribution) SampleIndex(n uint64) (uint64, error) { + if d.delegate == nil { + return 0, nil + } + return d.delegate.SampleIndex(n) +} + +func (d *Distribution) UnmarshalJSON(data []byte) error { + var temp struct { + Name string `json:"Name"` + } + if err := json.Unmarshal(data, &temp); err != nil { + return err + } + d.name = temp.Name + switch d.name { + case "": + return nil + case "uniform": + var uniform UniformDistribution + if err := json.Unmarshal(data, &uniform); err != nil { + return err + } + d.delegate = &uniform + return nil + case "zipfian": + var zipfian ZipfianDistribution + if err := json.Unmarshal(data, &zipfian); err != nil { + return err + } + if err := zipfian.validate(); err != nil { + return err + } + d.delegate = &zipfian + return nil + default: + return fmt.Errorf("unknown distribution name: %s", d.name) + } +} + +// UniformDistribution draws each index with equal probability. +type UniformDistribution struct{} + +func (UniformDistribution) SampleIndex(n uint64) (uint64, error) { + // PLT-460: implement the seeded uniform draw. Out of scope for PLT-455 + // (wire format + validation only). + return 0, nil +} + +// ZipfianDistribution draws indices with a Zipf-distributed skew controlled by +// theta. theta == 0 is uniform; larger theta concentrates draws on low indices. +type ZipfianDistribution struct { + Theta float64 `json:"theta"` +} + +// zipfianThetaMax bounds theta to the range over which the YCSB precomputed-zeta +// generator (PLT-460) is numerically well-behaved. +const zipfianThetaMax = 1.0 + +func (z *ZipfianDistribution) validate() error { + if z.Theta < 0 || z.Theta >= zipfianThetaMax { + return fmt.Errorf("zipfian theta out of range: %v (want [0, %v))", z.Theta, zipfianThetaMax) + } + return nil +} + +func (z *ZipfianDistribution) SampleIndex(n uint64) (uint64, error) { + // PLT-460: implement the YCSB precomputed-zeta zipfian draw with a seeded + // RNG. Out of scope for PLT-455 (wire format + validation only). + return 0, nil +} diff --git a/config/distribution_test.go b/config/distribution_test.go new file mode 100644 index 0000000..0dc32b2 --- /dev/null +++ b/config/distribution_test.go @@ -0,0 +1,68 @@ +package config_test + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/sei-protocol/sei-load/config" + "github.com/stretchr/testify/require" +) + +func TestDistribution(t *testing.T) { + t.Parallel() + t.Run("empty", func(t *testing.T) { + var subject config.Distribution + require.NoError(t, subject.UnmarshalJSON([]byte(`{}`))) + require.Empty(t, subject.Name()) + idx, err := subject.SampleIndex(100) + require.NoError(t, err) + require.Zero(t, idx) + }) + t.Run("uniform", func(t *testing.T) { + var subject config.Distribution + require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"uniform"}`))) + require.Equal(t, "uniform", subject.Name()) + }) + t.Run("zipfian", func(t *testing.T) { + var subject config.Distribution + require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":0.99}`))) + require.Equal(t, "zipfian", subject.Name()) + }) + t.Run("zipfian_theta_below_range", func(t *testing.T) { + var subject config.Distribution + require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":-0.1}`))) + }) + t.Run("zipfian_theta_above_range", func(t *testing.T) { + var subject config.Distribution + require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":1.0}`))) + }) + t.Run("unknown", func(t *testing.T) { + var subject config.Distribution + require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"weibull"}`))) + }) +} + +// TestScenarioDistributionAdditive proves the new fields are additive: a profile +// carrying no distribution fields parses unchanged and round-trips without +// introducing any distribution keys. +func TestScenarioDistributionAdditive(t *testing.T) { + t.Parallel() + path := filepath.Join("..", "profiles", "conflict.json") + original, err := os.ReadFile(path) + require.NoError(t, err) + + var cfg config.LoadConfig + require.NoError(t, json.Unmarshal(original, &cfg)) + + for _, s := range cfg.Scenarios { + require.Nil(t, s.KeyDistribution, "no distribution expected in baseline profile") + require.Nil(t, s.SizeDistribution, "no distribution expected in baseline profile") + } + + remarshaled, err := json.Marshal(cfg) + require.NoError(t, err) + require.NotContains(t, string(remarshaled), "keyDistribution") + require.NotContains(t, string(remarshaled), "sizeDistribution") +}