Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
enhance: enable to show etcd kv tree
Browse files Browse the repository at this point in the history
Signed-off-by: Wei Liu <[email protected]>
weiliu1031 committed Jan 15, 2025
1 parent 8694620 commit c2e0a90
Showing 1 changed file with 145 additions and 0 deletions.
145 changes: 145 additions & 0 deletions states/etcd/show/etcd_kv_tree.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
package show

import (
"context"
"fmt"
"sort"
"strconv"
"strings"

clientv3 "go.etcd.io/etcd/client/v3"

"github.com/milvus-io/birdwatcher/framework"
)

type EtcdKVTree struct {
framework.ParamBase `use:"show etcd-kv-tree" desc:"show etcd kv tree with key size of each prefix"`
Prefix string `name:"prefix" default:"" desc:"the kv prefix to show"`
Level string `name:"level" default:"1" desc:"the level of kv tree to show"`
TopK string `name:"topK" default:"10" desc:"the number of top prefixes to show per level"`
}

// EtcdKVTreeCommand retrieves and prints the top K prefixes and their key counts up to the specified level
func (c *ComponentShow) EtcdKVTreeCommand(ctx context.Context, p *EtcdKVTree) error {
// Parse the Level string to an integer
level, err := strconv.Atoi(p.Level)
if err != nil || level <= 0 {
return fmt.Errorf("invalid level: %s", p.Level)
}

// Parse the TopK string to an integer
topK, err := strconv.Atoi(p.TopK)
if err != nil || topK <= 0 {
topK = 10 // Default to 10 if parsing fails
}

// Fetch all keys under the given prefix
resp, err := c.client.Get(ctx, p.Prefix, clientv3.WithPrefix())
if err != nil {
return err
}

// Extract keys from the response
keys := make([]string, len(resp.Kvs))
for i, kv := range resp.Kvs {
keys[i] = string(kv.Key)
}

// Count keys for prefixes up to the specified level
result := countKeysAtEachLevel(keys, p.Prefix, level)

// Print the result with topK prefixes for each level in order
printLevelsInOrder(result, topK)

return nil
}

// countKeysAtEachLevel counts the keys for each prefix at each level up to the specified level
func countKeysAtEachLevel(keys []string, basePrefix string, maxLevel int) map[int]map[string]int {
levelStats := make(map[int]map[string]int)

for _, key := range keys {
// Ensure the key is under the base prefix
if !strings.HasPrefix(key, basePrefix) {
continue
}

// Process prefixes for each level up to maxLevel
for level := 1; level <= maxLevel; level++ {
prefix := getNthLevelPrefix(key, basePrefix, level)
if prefix == "" {
break
}

if _, exists := levelStats[level]; !exists {
levelStats[level] = make(map[string]int)
}
levelStats[level][prefix]++
}
}

return levelStats
}

// printLevelsInOrder ensures that levels are printed in order from 1 to maxLevel
func printLevelsInOrder(result map[int]map[string]int, topK int) {
// Iterate over the levels in sorted order (from 1 to maxLevel)
for level := 1; level <= len(result); level++ {
if prefixes, exists := result[level]; exists {
printTopKPrefixes(level, prefixes, topK)
}
}
}

// printTopKPrefixes prints the top K prefixes for a given level
func printTopKPrefixes(level int, prefixes map[string]int, topK int) {
// Convert map to slice for sorting
type prefixCount struct {
prefix string
count int
}
var sortedPrefixes []prefixCount
for prefix, count := range prefixes {
sortedPrefixes = append(sortedPrefixes, prefixCount{prefix, count})
}

// Sort the prefixes by count in descending order
sort.Slice(sortedPrefixes, func(i, j int) bool {
return sortedPrefixes[i].count > sortedPrefixes[j].count
})

// Keep only the top K prefixes
if len(sortedPrefixes) > topK {
sortedPrefixes = sortedPrefixes[:topK]
}

// Print the result
fmt.Printf("Level %d:\n", level)
for _, p := range sortedPrefixes {
fmt.Printf(" Prefix: %s, Key Count: %d\n", p.prefix, p.count)
}
}

// getNthLevelPrefix extracts the prefix up to the specified level under basePrefix
func getNthLevelPrefix(key, basePrefix string, level int) string {
if level <= 0 || !strings.HasPrefix(key, basePrefix) {
return ""
}

parts := strings.Split(strings.TrimPrefix(key, basePrefix), "/")
if len(parts) < level {
return ""
}

builder := strings.Builder{}
builder.WriteString(strings.Trim(basePrefix, "/"))

for i := 0; i < level; i++ {
if part := strings.Trim(parts[i], "/"); part != "" {
builder.WriteByte('/')
builder.WriteString(part)
}
}

return builder.String()
}

0 comments on commit c2e0a90

Please sign in to comment.