diff --git a/README.md b/README.md index c7fb17f..7a83ed4 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,7 @@ Options: 1. '1KB~1MB' get keys with size in range [1KB, 1MB] 2. '10MB~inf' magic variable 'inf' represents the Infinity 3. '1024~10KB' get keys with size in range [0Bytes, 10KB] + -sample sample rate between (0, 1] for sampling keys -concurrent The number of concurrent json converters. 4 by default. -show-global-meta Show global meta likes redis-verion/ctime/functions -no-expired filter expired keys(deprecated, please use 'expire' option) @@ -618,6 +619,19 @@ rdb -c json -o dump.json -size 1KB~1MB cases/memory.rdb rdb -c memory -o memory.csv -size 10MB~inf cases/memory.rdb ``` +# Sample Filter + +The `-sample` parameter can be configured to sample keys from large RDB files with a sampling rate in `(0, 1]`. + +Examples: +```bash +# Sample 10% of keys to generate JSON +rdb -c json -o dump.json -sample 0.1 dump.rdb + +# Sample 20% of keys to generate memory report +rdb -c memory -o memory.csv -sample 0.2 dump.rdb +``` + # Expiration Filter The `-expire` parameter can be configured to filter based on the expiration time. diff --git a/README_CN.md b/README_CN.md index 8886c4d..e3f4501 100644 --- a/README_CN.md +++ b/README_CN.md @@ -60,6 +60,7 @@ Options: 1. '1KB~1MB' get keys with size in range [1KB, 1MB] 2. '10MB~inf' magic variable 'inf' represents the Infinity 3. '1024~10KB' get keys with size in range [0Bytes, 10KB] + -sample sample rate between (0, 1] for sampling keys -concurrent The number of concurrent json converters. 4 by default. -show-global-meta Show global meta likes redis-verion/ctime/functions -no-expired filter expired keys(deprecated, please use 'expire' option) @@ -544,6 +545,19 @@ rdb -c json -o dump.json -size 1KB~1MB cases/memory.rdb rdb -c memory -o memory.csv -size 10MB~inf cases/memory.rdb ``` +# 抽样过滤器 + +`-sample` 参数可以配置对大型 RDB 文件中的键进行抽样处理,取值范围为 `(0, 1]`。 + +示例: +```bash +# 抽样 10% 的键生成 JSON +rdb -c json -o dump.json -sample 0.1 dump.rdb + +# 抽样 20% 的键生成内存报告 +rdb -c memory -o memory.csv -sample 0.2 dump.rdb +``` + # Expiration 过滤器 `-expire` 参数可以配置根据过期时间过滤. diff --git a/cmd.go b/cmd.go index 795f131..0a308f7 100644 --- a/cmd.go +++ b/cmd.go @@ -32,6 +32,7 @@ Options: 1. '1KB~1MB' get keys with size in range [1KB, 1MB] 2. '10MB~inf' magic variable 'inf' represents the Infinity 3. '1024~10KB' get keys with size in range [0Bytes, 10KB] + -sample sample rate between (0, 1] for sampling keys -concurrent The number of concurrent json converters. 4 by default. -show-global-meta Show global meta likes redis-verion/ctime/functions -no-expired filter expired keys(deprecated, please use 'expire' option) @@ -82,6 +83,7 @@ func main() { var concurrent int var showGlobalMeta bool var prefixSeps separators + var sample float64 var err error flagSet.StringVar(&cmd, "c", "", "command for rdb: json") flagSet.StringVar(&output, "o", "", "output file path") @@ -93,6 +95,7 @@ func main() { flagSet.StringVar(®exExpr, "regex", "", "regex expression") flagSet.StringVar(&expirationExpr, "expire", "", "expiration filter expression") flagSet.StringVar(&sizeExpr, "size", "", "size filter expression") + flagSet.Float64Var(&sample, "sample", 0, "sample rate between (0, 1] for sampling keys") flagSet.BoolVar(&noExpired, "no-expired", false, "filter expired keys(deprecated, please use expire)") flagSet.Var(&prefixSeps, "prefix-sep", "separator for prefix analysis (flat-map mode, constant memory)") flagSet.BoolVar(&showGlobalMeta, "show-global-meta", false, "Show global meta likes redis-verion/ctime/functions") @@ -124,6 +127,15 @@ func main() { if sizeExpr != "" { options = append(options, helper.WithSizeOption(sizeExpr)) } + var hasSample bool + flagSet.Visit(func(f *flag.Flag) { + if f.Name == "sample" { + hasSample = true + } + }) + if hasSample { + options = append(options, helper.WithSampleOption(sample)) + } if showGlobalMeta { options = append(options, helper.WithGlobalMeta()) } @@ -147,7 +159,7 @@ func main() { case "memory": err = helper.MemoryProfile(src, output, options...) case "aof": - err = helper.ToAOF(src, output, options) + err = helper.ToAOF(src, output, options...) case "bigkey": err = helper.FindBiggestKeys(src, n, outputFile, options...) case "hotkey": diff --git a/cmd_test.go b/cmd_test.go index 152c3e5..1331979 100644 --- a/cmd_test.go +++ b/cmd_test.go @@ -69,6 +69,32 @@ func TestCmd(t *testing.T) { t.Error("command prefix with prefix-sep failed") } + os.Args = []string{"", "-c", "memory", "-sample", "0.5", "-o", "tmp/memory_sample.csv", "cases/memory.rdb"} + main() + if f, _ := os.Stat("tmp/memory_sample.csv"); f == nil { + t.Error("command memory with sample failed") + } + + os.Args = []string{"", "-c", "json", "-sample", "0.5", "-o", "tmp/cmd_sample.json", "cases/memory.rdb"} + main() + if f, _ := os.Stat("tmp/cmd_sample.json"); f == nil { + t.Error("command json with sample failed") + } + + os.Args = []string{"", "-c", "bigkey", "-sample", "0.5", "-n", "10", "-o", "tmp/bigkey_sample.csv", "cases/memory.rdb"} + main() + if f, _ := os.Stat("tmp/bigkey_sample.csv"); f == nil { + t.Error("command bigkey with sample failed") + } + + // invalid sample rate commands + os.Args = []string{"", "-c", "memory", "-sample", "-0.1", "-o", "tmp/memory_invalid1.csv", "cases/memory.rdb"} + main() + os.Args = []string{"", "-c", "memory", "-sample", "0", "-o", "tmp/memory_invalid2.csv", "cases/memory.rdb"} + main() + os.Args = []string{"", "-c", "memory", "-sample", "1.5", "-o", "tmp/memory_invalid3.csv", "cases/memory.rdb"} + main() + // test error command line os.Args = []string{"", "-c", "json", "-o", "tmp/output", "/none/a"} main() diff --git a/helper/regex.go b/helper/regex.go index 8b58cd7..82083a1 100644 --- a/helper/regex.go +++ b/helper/regex.go @@ -211,8 +211,15 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) { var expirationOpt ExpirationOption var sizeOpt SizeOption var globalMetaOpt GlobalMetaOption + var sampleOpt *SampleOption for _, opt := range options { switch o := opt.(type) { + case []interface{}: + var err error + dec, err = wrapDecoder(dec, o...) + if err != nil { + return nil, err + } case RegexOption: regexOpt = o case NoExpiredOption: @@ -223,6 +230,14 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) { sizeOpt = o case GlobalMetaOption: globalMetaOpt = o + case SampleOption: + sampleOpt = &o + } + } + if sampleOpt != nil { + rate := float64(*sampleOpt) + if rate <= 0 || rate > 1 { + return nil, fmt.Errorf("illegal sample rate: %v, should be in range (0, 1]", rate) } } if globalMetaOpt { @@ -274,5 +289,11 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) { sizeRange: rng, } } + if sampleOpt != nil { + rate := float64(*sampleOpt) + if rate < 1.0 { + dec = newSampleDecoder(dec, rate) + } + } return dec, nil } diff --git a/helper/sample.go b/helper/sample.go new file mode 100644 index 0000000..5feb8f7 --- /dev/null +++ b/helper/sample.go @@ -0,0 +1,52 @@ +package helper + +import ( + "math/rand" + "time" + + "github.com/hdt3213/rdb/model" +) + +// SampleOption enables sampling keys with a rate in (0, 1] +type SampleOption float64 + +// WithSampleOption creates a SampleOption from sample rate +func WithSampleOption(rate float64) SampleOption { + return SampleOption(rate) +} + +// WithSample creates a SampleOption from sample rate +func WithSample(rate float64) SampleOption { + return SampleOption(rate) +} + +type sampleDecoder struct { + dec decoder + rate float64 + rnd *rand.Rand +} + +func newSampleDecoder(dec decoder, rate float64) *sampleDecoder { + return &sampleDecoder{ + dec: dec, + rate: rate, + rnd: rand.New(rand.NewSource(time.Now().UnixNano())), + } +} + +func newSampleDecoderWithRand(dec decoder, rate float64, rnd *rand.Rand) *sampleDecoder { + return &sampleDecoder{ + dec: dec, + rate: rate, + rnd: rnd, + } +} + +func (d *sampleDecoder) Parse(cb func(object model.RedisObject) bool) error { + return d.dec.Parse(func(object model.RedisObject) bool { + if d.rnd.Float64() < d.rate { + return cb(object) + } + return true + }) +} diff --git a/helper/sample_test.go b/helper/sample_test.go new file mode 100644 index 0000000..8079b5e --- /dev/null +++ b/helper/sample_test.go @@ -0,0 +1,231 @@ +package helper + +import ( + "math/rand" + "os" + "path/filepath" + "strconv" + "testing" + + "github.com/hdt3213/rdb/model" +) + +type mockDecoder struct { + objects []model.RedisObject +} + +func (m *mockDecoder) Parse(cb func(object model.RedisObject) bool) error { + for _, obj := range m.objects { + if !cb(obj) { + break + } + } + return nil +} + +func generateMockObjects(n int) []model.RedisObject { + objs := make([]model.RedisObject, n) + for i := 0; i < n; i++ { + objs[i] = &model.StringObject{ + BaseObject: &model.BaseObject{ + DB: 0, + Key: "key_" + strconv.Itoa(i), + }, + Value: []byte("val"), + } + } + return objs +} + +func TestSampleDecoderBasic(t *testing.T) { + total := 1000 + mock := &mockDecoder{objects: generateMockObjects(total)} + rnd := rand.New(rand.NewSource(42)) + sampler := newSampleDecoderWithRand(mock, 0.5, rnd) + + var sampled []model.RedisObject + err := sampler.Parse(func(object model.RedisObject) bool { + sampled = append(sampled, object) + return true + }) + if err != nil { + t.Fatalf("unexpected error during parse: %v", err) + } + + count := len(sampled) + // With seed 42, 1000 items and rate 0.5, count should be reasonably close to 500 (e.g. 400 to 600) + if count < 400 || count > 600 { + t.Fatalf("expected around 500 sampled items, got %d", count) + } +} + +func TestSampleDecoderEarlyTermination(t *testing.T) { + total := 100 + mock := &mockDecoder{objects: generateMockObjects(total)} + rnd := rand.New(rand.NewSource(42)) + sampler := newSampleDecoderWithRand(mock, 1.0, rnd) + + count := 0 + err := sampler.Parse(func(object model.RedisObject) bool { + count++ + return count < 5 // stop after 5 items + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if count != 5 { + t.Fatalf("expected early termination at 5 items, got %d", count) + } +} + +func TestWrapDecoderWithSampleOption(t *testing.T) { + mock := &mockDecoder{objects: generateMockObjects(10)} + + // Test invalid negative rate + _, err := wrapDecoder(mock, WithSampleOption(-0.5)) + if err == nil { + t.Fatal("expected error for negative sample rate") + } + + // Test invalid zero rate + _, err = wrapDecoder(mock, WithSampleOption(0)) + if err == nil { + t.Fatal("expected error for zero sample rate") + } + + // Test invalid rate > 1 + _, err = wrapDecoder(mock, WithSampleOption(1.5)) + if err == nil { + t.Fatal("expected error for sample rate > 1") + } + + // Test valid rate == 1.0 + dec, err := wrapDecoder(mock, WithSampleOption(1.0)) + if err != nil { + t.Fatalf("unexpected error for sample rate 1.0: %v", err) + } + var count int + err = dec.Parse(func(object model.RedisObject) bool { + count++ + return true + }) + if err != nil || count != 10 { + t.Fatalf("expected all 10 items for rate 1.0, got %d (err: %v)", count, err) + } + + // Test WithSample alias function + dec2, err := wrapDecoder(mock, WithSample(0.5)) + if err != nil { + t.Fatalf("unexpected error for WithSample(0.5): %v", err) + } + if dec2 == nil { + t.Fatal("expected non-nil decoder") + } +} + +func TestSampleIntegrationWithRDB(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "rdb_sample_test_*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer func() { + _ = os.RemoveAll(tmpDir) + }() + + srcRdb := filepath.Join("../cases", "memory.rdb") + + // 1. Test MemoryProfile with sampling + memOut := filepath.Join(tmpDir, "memory_sampled.csv") + err = MemoryProfile(srcRdb, memOut, WithSampleOption(1.0)) + if err != nil { + t.Fatalf("MemoryProfile with sample 1.0 failed: %v", err) + } + if fi, err := os.Stat(memOut); err != nil || fi.Size() == 0 { + t.Fatalf("expected non-empty memory csv output, got err=%v", err) + } + + // 2. Test FindBiggestKeys with sampling + bigkeyFile, err := os.Create(filepath.Join(tmpDir, "bigkey_sampled.csv")) + if err != nil { + t.Fatalf("failed to create bigkey file: %v", err) + } + defer func() { + _ = bigkeyFile.Close() + }() + err = FindBiggestKeys(srcRdb, 5, bigkeyFile, WithSampleOption(1.0)) + if err != nil { + t.Fatalf("FindBiggestKeys with sample 1.0 failed: %v", err) + } + + // 3. Test ToJsons with sampling + jsonOut := filepath.Join(tmpDir, "sampled.json") + err = ToJsons(srcRdb, jsonOut, WithSampleOption(1.0), WithConcurrent(1)) + if err != nil { + t.Fatalf("ToJsons with sample 1.0 failed: %v", err) + } + + // 4. Test invalid sample rate propagation + err = MemoryProfile(srcRdb, filepath.Join(tmpDir, "invalid.csv"), WithSampleOption(-0.1)) + if err == nil { + t.Fatal("expected error from MemoryProfile with invalid sample rate") + } +} + +func TestSampleIntegrationToAOF(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "rdb_aof_sample_test_*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer func() { + _ = os.RemoveAll(tmpDir) + }() + + srcRdb := filepath.Join("../cases", "memory.rdb") + aofOut := filepath.Join(tmpDir, "sampled.aof") + err = ToAOF(srcRdb, aofOut, WithSampleOption(1.0)) + if err != nil { + t.Fatalf("ToAOF with sample 1.0 failed: %v", err) + } + if fi, err := os.Stat(aofOut); err != nil || fi.Size() == 0 { + t.Fatalf("expected non-empty aof output, got err=%v", err) + } +} + +func TestSampleIntegrationPrefix(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "rdb_prefix_sample_test_*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + defer func() { + _ = os.RemoveAll(tmpDir) + }() + + srcRdb := filepath.Join("../cases", "tree.rdb") + + sepOut, err := os.Create(filepath.Join(tmpDir, "sep_prefix.csv")) + if err != nil { + t.Fatalf("failed to create sep prefix file: %v", err) + } + err = SepPrefixAnalyse(srcRdb, 10, 3, []string{":"}, sepOut, WithSampleOption(1.0)) + _ = sepOut.Close() + if err != nil { + t.Fatalf("SepPrefixAnalyse with sample 1.0 failed: %v", err) + } + if fi, err := os.Stat(sepOut.Name()); err != nil || fi.Size() == 0 { + t.Fatal("expected non-empty SepPrefixAnalyse output") + } + + prefixOut, err := os.Create(filepath.Join(tmpDir, "prefix.csv")) + if err != nil { + t.Fatalf("failed to create prefix file: %v", err) + } + err = PrefixAnalyse(srcRdb, 10, 3, prefixOut, WithSampleOption(1.0)) + _ = prefixOut.Close() + if err != nil { + t.Fatalf("PrefixAnalyse with sample 1.0 failed: %v", err) + } + if fi, err := os.Stat(prefixOut.Name()); err != nil || fi.Size() == 0 { + t.Fatal("expected non-empty PrefixAnalyse output") + } +} +