Skip to content

Commit

Permalink
improve loadtest watch cmd
Browse files Browse the repository at this point in the history
  • Loading branch information
fspmarshall committed Jul 30, 2024
1 parent 558d71a commit 65f003e
Showing 1 changed file with 61 additions and 9 deletions.
70 changes: 61 additions & 9 deletions tool/tctl/common/loadtest_command.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import (

"github.com/alecthomas/kingpin/v2"
"github.com/google/uuid"
"github.com/gravitational/teleport"
"github.com/gravitational/trace"
log "github.com/sirupsen/logrus"

Expand All @@ -39,6 +40,7 @@ import (
"github.com/gravitational/teleport/lib/cache"
"github.com/gravitational/teleport/lib/service/servicecfg"
"github.com/gravitational/teleport/lib/services"
"github.com/gravitational/teleport/lib/utils"
)

// LoadtestCommand implements the `tctl loadtest` family of commands.
Expand All @@ -56,7 +58,9 @@ type LoadtestCommand struct {
ttl time.Duration
concurrency int

kind string
kind string
ops string
format string
}

// Initialize allows LoadtestCommand to plug itself into the CLI parser
Expand All @@ -76,6 +80,8 @@ func (c *LoadtestCommand) Initialize(app *kingpin.Application, config *servicecf

c.watch = loadtest.Command("watch", "Monitor event stream").Hidden()
c.watch.Flag("kind", "Resource kind(s) to watch, e.g. --kind=node,user,role").StringVar(&c.kind)
c.watch.Flag("ops", "Operations to watch, e.g. --ops=put,del").Default("put,del").StringVar(&c.ops)
c.watch.Flag("format", "Output format").Default(teleport.Text).StringVar(&c.format)
}

// TryRun takes the CLI command as an argument (like "loadtest node-heartbeats") and executes it.
Expand Down Expand Up @@ -224,6 +230,27 @@ func (c *LoadtestCommand) Watch(ctx context.Context, client *authclient.Client)
})
}

ops := make(map[types.OpType]struct{})
for _, op := range strings.Split(c.ops, ",") {
op = strings.TrimSpace(op)
if op == "" {
continue
}

switch op {
case "put":
ops[types.OpPut] = struct{}{}
case "del":
ops[types.OpDelete] = struct{}{}
default:
return trace.BadParameter("unknown operation: %v", op)
}
}

if len(ops) == 0 {
return trace.BadParameter("no operations specified")
}

var allowPartialSuccess bool
if len(kinds) == 0 {
// use auth watch kinds by default
Expand Down Expand Up @@ -254,19 +281,26 @@ func (c *LoadtestCommand) Watch(ctx context.Context, client *authclient.Client)
skinds = append(skinds, k.Kind)
}

fmt.Printf("INIT: %v\n", skinds)
fmt.Fprintf(os.Stderr, "INIT: %v\n", skinds)
case <-watcher.Done():
return trace.Errorf("failed to get init event: %v", watcher.Error())
}

for {
select {
case event := <-watcher.Events():
if _, ok := ops[event.Type]; !ok {
continue
}
switch event.Type {
case types.OpPut:
printEvent("PUT", event.Resource)
if err := printEvent("PUT", event.Resource, c.format); err != nil {
return trace.Wrap(err)
}
case types.OpDelete:
printEvent("DEL", event.Resource)
if err := printEvent("DEL", event.Resource, c.format); err != nil {
return trace.Wrap(err)
}
default:
return trace.BadParameter("expected put or del event, got %v instead", event.Type)
}
Expand All @@ -280,10 +314,28 @@ func (c *LoadtestCommand) Watch(ctx context.Context, client *authclient.Client)
}
}

func printEvent(ekind string, rsc types.Resource) {
if sk := rsc.GetSubKind(); sk != "" {
fmt.Printf("%s: %s/%s/%s\n", ekind, rsc.GetKind(), sk, rsc.GetName())
} else {
fmt.Printf("%s: %s/%s\n", ekind, rsc.GetKind(), rsc.GetName())
func printEvent(ekind string, rsc types.Resource, format string) error {
var ln string
switch format {
case teleport.Text:
if sk := rsc.GetSubKind(); sk != "" {
ln = fmt.Sprintf("%s: %s/%s/%s", ekind, rsc.GetKind(), sk, rsc.GetName())
} else {
ln = fmt.Sprintf("%s: %s/%s", ekind, rsc.GetKind(), rsc.GetName())
}
case teleport.JSON:
b, err := utils.FastMarshal(map[string]any{
"op": ekind,
"resource": rsc,
})
if err != nil {
return trace.Wrap(err)
}
ln = string(b)
default:
return trace.BadParameter("unknown format: %v", format)
}

fmt.Println(ln)
return nil
}

0 comments on commit 65f003e

Please sign in to comment.