vpp-maglev/cmd/client/watch.go

// SPDX-License-Identifier: Apache-2.0

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"

	"git.ipng.ch/ipng/golang-cli/keypress"
	"git.ipng.ch/ipng/vpp-maglev/internal/grpcapi"
)

func dynWatchEventOpts(_ context.Context, _ grpcapi.MaglevClient, _ []string) []string {
	return []string{"num", "log", "backend", "frontend"}
}

// runWatchEvents implements 'watch events [num <n>] [log [level <level>]] [backend] [frontend]'.
// All tokens after 'events' are captured as args by the circular slot node in the tree.
// If none of log/backend/frontend are mentioned, all three default to true.
func runWatchEvents(ctx context.Context, client grpcapi.MaglevClient, args []string) error {
	var maxEvents int // 0 = unlimited
	var wantLog, wantBackend, wantFrontend bool
	logLevel := ""
	anyExplicit := false

	for i := 0; i < len(args); {
		switch args[i] {
		case "num":
			if i+1 >= len(args) {
				return fmt.Errorf("num requires a count argument")
			}
			n, err := strconv.Atoi(args[i+1])
			if err != nil || n < 1 {
				return fmt.Errorf("num: invalid count %q", args[i+1])
			}
			maxEvents = n
			i += 2
		case "log":
			wantLog = true
			anyExplicit = true
			if i+1 < len(args) && args[i+1] == "level" {
				if i+2 >= len(args) {
					return fmt.Errorf("log level requires a level argument")
				}
				logLevel = args[i+2]
				i += 3
			} else {
				i++
			}
		case "backend":
			wantBackend = true
			anyExplicit = true
			i++
		case "frontend":
			wantFrontend = true
			anyExplicit = true
			i++
		default:
			return fmt.Errorf("unknown watch option %q; expected: num, log, backend, frontend", args[i])
		}
	}

	if !anyExplicit {
		wantLog, wantBackend, wantFrontend = true, true, true
	}

	boolp := func(b bool) *bool { v := b; return &v }
	req := &grpcapi.WatchRequest{
		Log:      boolp(wantLog),
		LogLevel: logLevel,
		Backend:  boolp(wantBackend),
		Frontend: boolp(wantFrontend),
	}

	// Cancel the stream on keypress or signal.
	watchCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	go keypress.WaitForKey(watchCtx, cancel)

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	defer signal.Stop(sigCh)
	go func() {
		select {
		case <-sigCh:
			cancel()
		case <-watchCtx.Done():
		}
	}()

	stream, err := client.WatchEvents(watchCtx, req)
	if err != nil {
		return err
	}

	marshaler := protojson.MarshalOptions{}
	count := 0
	for {
		ev, err := stream.Recv()
		if err != nil {
			if watchCtx.Err() != nil {
				return nil // stopped by keypress or signal
			}
			if st, ok := status.FromError(err); ok && st.Code() == codes.Canceled {
				return nil
			}
			return err
		}
		data, err := marshaler.Marshal(ev)
		if err != nil {
			return fmt.Errorf("marshal event: %w", err)
		}
		fmt.Printf("%s\n", data)
		count++
		if maxEvents > 0 && count >= maxEvents {
			return nil
		}
	}
}

// Stopping the watch on a keypress now uses keypress.WaitForKey from the
// golang-cli library (per-platform cbreak handling, with a non-tty fallback).