flamegraph support multi separators

README.md

@@ -140,15 +140,19 @@ aaaaaaa
 
 In many cases there is not a few very large key but lots of small keys that occupied most memory.
 
-RDB tool could separate keys by given delimeter, then aggregate keys with same prefix.
+RDB tool could separate keys by the given delimeters, then aggregate keys with same prefix.
 
 Finally RDB tool presents the result as flame graph, with which you could find out which kind of keys consumed most
 memory.
 
+![](https://s2.loli.net/2022/03/27/eNGvVIdAuWp8EhT.png)
+
+In this example, the keys of pattern `Comment:*` use 8.463% memory.
+
 Usage:
 
 ```
-rdb -c flamegraph [-port <port>] [-sep <separator>] <source_path>
+rdb -c flamegraph [-port <port>] [-sep <separator1>] [-sep <separator2>] <source_path>
 ```
 
 Example:
@@ -157,10 +161,6 @@ Example:
 rdb -c flamegraph -port 16379 -sep : dump.rdb
 ```
 
-![](https://s2.loli.net/2022/03/27/eNGvVIdAuWp8EhT.png)
-
-In the example, the keys of pattern `Comment:*` use 8.463% memory.
-
 # Customize data usage
 
 ```go

cmd.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/hdt3213/rdb/helper"
 	"os"
+	"strings"
 )
 
 const help = `
@@ -14,7 +15,8 @@ Options:
   -o output file path
   -n number of result 
   -port listen port for flame graph web service
-  -sep separator for flamegraph, rdb will separate key by it, default value is ":"
+  -sep separator for flamegraph, rdb will separate key by it, default value is ":". 
+		supporting multi separators: -sep sep1 -sep sep2 
 
 Examples:
 1. convert rdb to json
@@ -29,18 +31,29 @@ Examples:
   rdb -c flamegraph -port 16379 -sep : dump.rdb
 `
 
+type separators []string
+
+func (s *separators) String() string {
+	return strings.Join(*s, " ")
+}
+
+func (s *separators) Set(value string) error {
+	*s = append(*s, value)
+	return nil
+}
+
 func main() {
 	flagSet := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
 	var cmd string
 	var output string
 	var n int
 	var port int
-	var separator string
+	var seps separators
 	flagSet.StringVar(&cmd, "c", "", "command for rdb: json")
 	flagSet.StringVar(&output, "o", "", "output file path")
 	flagSet.IntVar(&n, "n", 0, "")
 	flagSet.IntVar(&port, "port", 0, "listen port for web")
-	flagSet.StringVar(&separator, "sep", "", "")
+	flagSet.Var(&seps, "sep", "separator for flamegraph")
 	_ = flagSet.Parse(os.Args[1:]) // ExitOnError
 	src := flagSet.Arg(0)
 
@@ -75,7 +88,7 @@ func main() {
 			err = helper.FindBiggestKeys(src, n, outputFile)
 		}
 	case "flamegraph":
-		_, err = helper.FlameGraph(src, port, separator)
+		_, err = helper.FlameGraph(src, port, seps...)
 		<-make(chan struct{})
 	default:
 		println("unknown command")

helper/flamegraph.go

@@ -16,13 +16,10 @@ import (
 var TrimThreshold = 1000
 
 // FlameGraph draws flamegraph in web page to analysis memory usage pattern
-func FlameGraph(rdbFilename string, port int, separator string) (chan<- struct{}, error) {
+func FlameGraph(rdbFilename string, port int, separator ...string) (chan<- struct{}, error) {
 	if rdbFilename == "" {
 		return nil, errors.New("src file path is required")
 	}
-	if separator == "" {
-		separator = ":"
-	}
 	if port == 0 {
 		port = 16379 // default port
 	}
@@ -56,9 +53,20 @@ func FlameGraph(rdbFilename string, port int, separator string) (chan<- struct{}
 	return d3flame.Web(data, port), nil
 }
 
-func addObject(root *d3flame.FlameItem, separator string, object model.RedisObject) {
+func split(s string, separators []string) []string {
+	sep := ":"
+	if len(separators) > 0 {
+		sep = separators[0]
+	}
+	for i := 1; i < len(separators); i++ {
+		s = strings.ReplaceAll(s, separators[i], sep)
+	}
+	return strings.Split(s, sep)
+}
+
+func addObject(root *d3flame.FlameItem, separators []string, object model.RedisObject) {
 	node := root
-	parts := strings.Split(object.GetKey(), separator)
+	parts := split(object.GetKey(), separators)
 	parts = append([]string{"db:" + strconv.Itoa(object.GetDBIndex())}, parts...)
 	for _, part := range parts {
 		if node.Children[part] == nil {

helper/flamegraph_test.go

@@ -0,0 +1,18 @@
+package helper
+
+import "testing"
+
+func TestSplit(t *testing.T) {
+	result := split("a:b:c", nil)
+	if len(result) != 3 || result[0] != "a" || result[1] != "b" || result[2] != "c" {
+		t.Errorf("wrong result: %+v", result)
+	}
+	result = split("a.b.c", []string{"."})
+	if len(result) != 3 || result[0] != "a" || result[1] != "b" || result[2] != "c" {
+		t.Errorf("wrong result: %+v", result)
+	}
+	result = split("a++b--c", []string{"++", "--"})
+	if len(result) != 3 || result[0] != "a" || result[1] != "b" || result[2] != "c" {
+		t.Errorf("wrong result: %+v", result)
+	}
+}