frostfs-node/cmd/frostfs-cli/modules/container/policy_playground.go

179 lines
4.8 KiB
Go
Raw Normal View History

package container
import (
"bufio"
"encoding/hex"
"fmt"
"io"
"os"
"strings"
internalclient "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/client"
"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/commonflags"
"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/key"
commonCmd "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/internal/common"
"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/netmap"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)
type policyPlaygroundREPL struct {
cmd *cobra.Command
args []string
nodes map[string]netmap.NodeInfo
}
func newPolicyPlaygroundREPL(cmd *cobra.Command, args []string) (*policyPlaygroundREPL, error) {
return &policyPlaygroundREPL{
cmd: cmd,
args: args,
nodes: map[string]netmap.NodeInfo{},
}, nil
}
func (repl *policyPlaygroundREPL) handleLs(args []string) error {
if len(args) > 0 {
return fmt.Errorf("too many arguments for command 'ls': got %d, want 0", len(args))
}
i := 1
for id, node := range repl.nodes {
var attrs []string
node.IterateAttributes(func(k, v string) {
attrs = append(attrs, fmt.Sprintf("%s:%s", k, v))
})
fmt.Printf("\t%2d: id=%s attrs={%v}\n", i, id, strings.Join(attrs, " "))
i++
}
return nil
}
func (repl *policyPlaygroundREPL) handleAdd(args []string) error {
if len(args) == 0 {
return fmt.Errorf("too few arguments for command 'add': got %d, want >0", len(args))
}
id := args[0]
key, err := hex.DecodeString(id)
if err != nil {
return fmt.Errorf("node id must be a hex string: got %q: %v", id, err)
}
node := repl.nodes[id]
node.SetPublicKey(key)
for _, attr := range args[1:] {
kv := strings.Split(attr, ":")
if len(kv) != 2 {
return fmt.Errorf("node attributes must be in the format 'KEY:VALUE': got %q", attr)
}
node.SetAttribute(kv[0], kv[1])
}
repl.nodes[id] = node
return nil
}
func (repl *policyPlaygroundREPL) handleRemove(args []string) error {
if len(args) == 0 {
return fmt.Errorf("too few arguments for command 'remove': got %d, want >0", len(args))
}
id := args[0]
if _, exists := repl.nodes[id]; exists {
delete(repl.nodes, id)
return nil
}
return fmt.Errorf("node not found: id=%q", id)
}
func (repl *policyPlaygroundREPL) handleEval(args []string) error {
policyStr := strings.Join(args, " ")
placementPolicy, err := parseContainerPolicy(repl.cmd, policyStr)
if err != nil {
return fmt.Errorf("parsing placement policy: %v", err)
}
nm := repl.netMap()
nodes, err := nm.ContainerNodes(*placementPolicy, nil)
if err != nil {
return fmt.Errorf("building container nodes: %v", err)
}
for i, ns := range nodes {
var ids []string
for _, node := range ns {
ids = append(ids, hex.EncodeToString(node.PublicKey()))
}
fmt.Printf("\t%2d: %v\n", i+1, ids)
}
return nil
}
func (repl *policyPlaygroundREPL) netMap() netmap.NetMap {
var nm netmap.NetMap
var nodes []netmap.NodeInfo
for _, node := range repl.nodes {
nodes = append(nodes, node)
}
nm.SetNodes(nodes)
return nm
}
func (repl *policyPlaygroundREPL) run() error {
if len(viper.GetString(commonflags.RPC)) > 0 {
key := key.GetOrGenerate(repl.cmd)
cli := internalclient.GetSDKClientByFlag(repl.cmd, key, commonflags.RPC)
var prm internalclient.NetMapSnapshotPrm
prm.SetClient(cli)
resp, err := internalclient.NetMapSnapshot(repl.cmd.Context(), prm)
commonCmd.ExitOnErr(repl.cmd, "unable to get netmap snapshot to populate initial netmap: %w", err)
for _, node := range resp.NetMap().Nodes() {
id := hex.EncodeToString(node.PublicKey())
repl.nodes[id] = node
}
}
cmdHandlers := map[string]func([]string) error{
"ls": repl.handleLs,
"add": repl.handleAdd,
"remove": repl.handleRemove,
"eval": repl.handleEval,
}
for reader := bufio.NewReader(os.Stdin); ; {
fmt.Print("> ")
line, err := reader.ReadString('\n')
if err != nil {
if err == io.EOF {
return nil
}
return fmt.Errorf("reading line: %v", err)
}
parts := strings.Fields(line)
if len(parts) == 0 {
continue
}
cmd := parts[0]
handler, exists := cmdHandlers[cmd]
if exists {
if err := handler(parts[1:]); err != nil {
fmt.Printf("error: %v\n", err)
}
} else {
fmt.Printf("error: unknown command %q\n", cmd)
}
}
}
var policyPlaygroundCmd = &cobra.Command{
Use: "policy-playground",
Short: "A REPL for testing placement policies",
Long: `A REPL for testing placement policies.
If a wallet and endpoint is provided, the initial netmap data will be loaded from the snapshot of the node. Otherwise, an empty playground is created.`,
Run: func(cmd *cobra.Command, args []string) {
repl, err := newPolicyPlaygroundREPL(cmd, args)
commonCmd.ExitOnErr(cmd, "could not create policy playground: %w", err)
commonCmd.ExitOnErr(cmd, "policy playground failed: %w", repl.run())
},
}
func initContainerPolicyPlaygroundCmd() {
commonflags.Init(policyPlaygroundCmd)
}