forked from TrueCloudLab/frostfs-node
258 lines
6.9 KiB
Go
258 lines
6.9 KiB
Go
package container
|
|
|
|
import (
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"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/chzyer/readline"
|
|
"github.com/spf13/cobra"
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
type policyPlaygroundREPL struct {
|
|
cmd *cobra.Command
|
|
nodes map[string]netmap.NodeInfo
|
|
}
|
|
|
|
func newPolicyPlaygroundREPL(cmd *cobra.Command) (*policyPlaygroundREPL, error) {
|
|
return &policyPlaygroundREPL{
|
|
cmd: cmd,
|
|
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:%q", 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) handleLoad(args []string) error {
|
|
if len(args) != 1 {
|
|
return fmt.Errorf("too few arguments for command 'add': got %d, want 1", len(args))
|
|
}
|
|
|
|
jsonNetmap := map[string]map[string]string{}
|
|
|
|
b, err := os.ReadFile(args[0])
|
|
if err != nil {
|
|
return fmt.Errorf("reading netmap file %q: %v", args[0], err)
|
|
}
|
|
|
|
if err := json.Unmarshal(b, &jsonNetmap); err != nil {
|
|
return fmt.Errorf("decoding json netmap: %v", err)
|
|
}
|
|
|
|
repl.nodes = make(map[string]netmap.NodeInfo)
|
|
for id, attrs := range jsonNetmap {
|
|
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 k, v := range attrs {
|
|
node.SetAttribute(k, v)
|
|
}
|
|
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.TrimSpace(strings.Join(args, " "))
|
|
var nodes [][]netmap.NodeInfo
|
|
nm := repl.netMap()
|
|
|
|
if strings.HasPrefix(policyStr, "CBF") || strings.HasPrefix(policyStr, "SELECT") || strings.HasPrefix(policyStr, "FILTER") {
|
|
// Assume that the input is a partial SELECT-FILTER expression.
|
|
// Full inline policies always start with UNIQUE or REP keywords,
|
|
// or different prefixes when it's the case of an external file.
|
|
sfExpr, err := netmap.DecodeSelectFilterString(policyStr)
|
|
if err != nil {
|
|
return fmt.Errorf("parsing select-filter expression: %v", err)
|
|
}
|
|
nodes, err = nm.SelectFilterNodes(sfExpr)
|
|
if err != nil {
|
|
return fmt.Errorf("building select-filter nodes: %v", err)
|
|
}
|
|
} else {
|
|
// Assume that the input is a full policy or input file otherwise.
|
|
placementPolicy, err := parseContainerPolicy(repl.cmd, policyStr)
|
|
if err != nil {
|
|
return fmt.Errorf("parsing placement policy: %v", err)
|
|
}
|
|
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
|
|
}
|
|
|
|
var policyPlaygroundCompleter = readline.NewPrefixCompleter(
|
|
readline.PcItem("list"),
|
|
readline.PcItem("ls"),
|
|
readline.PcItem("add"),
|
|
readline.PcItem("load"),
|
|
readline.PcItem("remove"),
|
|
readline.PcItem("rm"),
|
|
readline.PcItem("eval"),
|
|
)
|
|
|
|
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{
|
|
"list": repl.handleLs,
|
|
"ls": repl.handleLs,
|
|
"add": repl.handleAdd,
|
|
"load": repl.handleLoad,
|
|
"remove": repl.handleRemove,
|
|
"rm": repl.handleRemove,
|
|
"eval": repl.handleEval,
|
|
}
|
|
|
|
rl, err := readline.NewEx(&readline.Config{
|
|
Prompt: "> ",
|
|
InterruptPrompt: "^C",
|
|
AutoComplete: policyPlaygroundCompleter,
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("error initializing readline: %w", err)
|
|
}
|
|
defer rl.Close()
|
|
|
|
var exit bool
|
|
for {
|
|
line, err := rl.Readline()
|
|
if err != nil {
|
|
if errors.Is(err, readline.ErrInterrupt) {
|
|
if exit {
|
|
return nil
|
|
}
|
|
exit = true
|
|
continue
|
|
}
|
|
return fmt.Errorf("reading line: %v", err)
|
|
}
|
|
exit = false
|
|
|
|
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, _ []string) {
|
|
repl, err := newPolicyPlaygroundREPL(cmd)
|
|
commonCmd.ExitOnErr(cmd, "could not create policy playground: %w", err)
|
|
commonCmd.ExitOnErr(cmd, "policy playground failed: %w", repl.run())
|
|
},
|
|
}
|
|
|
|
func initContainerPolicyPlaygroundCmd() {
|
|
commonflags.Init(policyPlaygroundCmd)
|
|
}
|