feat: add Unix control service and rvc queue path

2026-09-06 09:57:36 +00:00
parent ca26d35632
commit 3f6fb89a41
10 changed files with 1371 additions and 5 deletions
@@ -13,10 +13,13 @@ import (
 	"os/signal"
 	"syscall"

+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
 	"github.com/rvbox/rvbox/internal/agentproto"
 	"github.com/rvbox/rvbox/internal/config"
+	"github.com/rvbox/rvbox/internal/server/control"
 	"github.com/rvbox/rvbox/internal/server/session"
 	"github.com/rvbox/rvbox/internal/server/store"
+	"google.golang.org/grpc"
 )

 func main() {
@@ -60,6 +63,21 @@ func run(configPath string) error {
 	if err != nil {
 		return fmt.Errorf("listen for agents: %w", err)
 	}
+	defer listener.Close()
+	controlListener, cleanupControl, err := control.ListenUnix(configured.Server.ControlSocket)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = cleanupControl() }()
+	controlService, err := control.NewService(control.Options{
+		Store: persistence, DefaultQueueTTL: configured.Queue.DefaultTTL, DefaultQueueTTLSet: true,
+		Limits: agentproto.Limits{MaxEnvelopeBytes: configured.Protocol.MaxAgentEnvelopeBytes, MaxExecutionSpecBytes: configured.Protocol.MaxExecutionSpecBytes, MaxRawChunkBytes: configured.Protocol.MaxRawChunkBytes, MaxScriptBytes: configured.Protocol.MaxScriptBytes, MaxDetailBytes: configured.Storage.ProtocolDetailMaxBytes},
+	})
+	if err != nil {
+		return err
+	}
+	grpcServer := grpc.NewServer(grpc.MaxRecvMsgSize(int(configured.Protocol.MaxControlRequestBytes)), grpc.MaxSendMsgSize(int(configured.Protocol.MaxControlRequestBytes)))
+	rvboxv1.RegisterControlServer(grpcServer, controlService)
 	agent := &session.AgentServer{
 		Store: persistence, Registry: session.NewRegistry(), Path: configured.Server.AgentPath,
 		Limits: agentproto.Limits{
@@ -73,8 +91,9 @@ func run(configPath string) error {
 		HeartbeatIdle: configured.Protocol.HeartbeatIdle, LivenessTimeout: configured.Protocol.LivenessTimeout,
 	}
 	httpServer := &http.Server{Handler: agent, ReadHeaderTimeout: configured.Flow.WriteDeadline}
-	serveError := make(chan error, 1)
+	serveError := make(chan error, 2)
 	go func() { serveError <- httpServer.Serve(listener) }()
+	go func() { serveError <- grpcServer.Serve(controlListener) }()

 	signals := make(chan os.Signal, 1)
 	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)
@@ -82,12 +101,22 @@ func run(configPath string) error {
 	select {
 	case err := <-serveError:
 		if errors.Is(err, http.ErrServerClosed) {
+			grpcServer.Stop()
 			return nil
 		}
+		grpcServer.Stop()
 		return err
 	case <-signals:
 		shutdownContext, cancel := context.WithTimeout(context.Background(), configured.Server.ShutdownGrace)
 		defer cancel()
-		return httpServer.Shutdown(shutdownContext)
+		httpErr := httpServer.Shutdown(shutdownContext)
+		grpcDone := make(chan struct{})
+		go func() { grpcServer.GracefulStop(); close(grpcDone) }()
+		select {
+		case <-grpcDone:
+		case <-shutdownContext.Done():
+			grpcServer.Stop()
+		}
+		return httpErr
 	}
 }
@@ -1,4 +1,287 @@
-// Command rvc is the RVBox control-plane command-line client.
+// Command rvc is the RVBox local control-plane CLI. It never opens the server
+// database; all state changes and reads go through the Unix gRPC socket.
 package main

-func main() {}
+import (
+	"context"
+	"crypto/sha256"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"net"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
+	"github.com/rvbox/rvbox/internal/domain"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const defaultControlSocket = "/run/rvbox/server.sock"
+
+func main() {
+	if err := run(os.Args[1:], os.Stdout, os.Stderr); err != nil {
+		fmt.Fprintln(os.Stderr, "rvc:", err)
+		os.Exit(1)
+	}
+}
+
+func run(args []string, output, diagnostics io.Writer) error {
+	if len(args) == 0 {
+		return errors.New("a command is required (stat or run)")
+	}
+	socket, args, err := globalSocket(args)
+	if err != nil {
+		return err
+	}
+	if len(args) == 0 {
+		return errors.New("a command is required (stat or run)")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	connection, err := dial(ctx, socket)
+	if err != nil {
+		return err
+	}
+	defer connection.Close()
+	client := rvboxv1.NewControlClient(connection)
+	switch args[0] {
+	case "stat":
+		return stat(ctx, client, args[1:], output)
+	case "run":
+		return runCommand(ctx, client, args[1:], output, diagnostics)
+	default:
+		return fmt.Errorf("unknown command %q", args[0])
+	}
+}
+
+func globalSocket(args []string) (string, []string, error) {
+	socket := defaultControlSocket
+	remaining := make([]string, 0, len(args))
+	for index := 0; index < len(args); index++ {
+		if args[index] == "--socket" {
+			if index+1 >= len(args) {
+				return "", nil, errors.New("--socket requires a path")
+			}
+			socket = args[index+1]
+			index++
+			continue
+		}
+		if strings.HasPrefix(args[index], "--socket=") {
+			socket = strings.TrimPrefix(args[index], "--socket=")
+			continue
+		}
+		remaining = append(remaining, args[index])
+	}
+	if socket == "" || !filepath.IsAbs(socket) {
+		return "", nil, errors.New("--socket must be an absolute path")
+	}
+	return socket, remaining, nil
+}
+
+func dial(ctx context.Context, socket string) (*grpc.ClientConn, error) {
+	dialer := func(ctx context.Context, _ string) (net.Conn, error) {
+		return (&net.Dialer{}).DialContext(ctx, "unix", socket)
+	}
+	return grpc.DialContext(ctx, "passthrough:///rvbox-control", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(dialer), grpc.WithBlock())
+}
+
+func stat(ctx context.Context, client rvboxv1.ControlClient, args []string, output io.Writer) error {
+	flags := flag.NewFlagSet("stat", flag.ContinueOnError)
+	flags.SetOutput(io.Discard)
+	all := flags.Bool("all", false, "show all command pages")
+	pageSize := flags.Uint("per-page", 100, "number of records per page")
+	if err := flags.Parse(args); err != nil {
+		return err
+	}
+	positionals := flags.Args()
+	if len(positionals) == 0 {
+		var token string
+		for {
+			response, err := client.ListClients(ctx, &rvboxv1.ListClientsRequest{PageSize: uint32(*pageSize), PageToken: token})
+			if err != nil {
+				return err
+			}
+			for _, item := range response.GetClients() {
+				fmt.Fprintf(output, "client %s connected=%t platform=%s running=%d queued=%d\n", item.GetClientId(), item.GetConnected(), item.GetPlatform(), item.GetRunningCommands(), item.GetQueuedCommands())
+			}
+			if !*all || response.GetNextPageToken() == "" {
+				return nil
+			}
+			token = response.GetNextPageToken()
+		}
+	}
+	if len(positionals) > 2 {
+		return errors.New("stat accepts at most CLIENT and ISSUE_UUID")
+	}
+	clientID := positionals[0]
+	if len(positionals) == 1 {
+		response, err := client.GetClient(ctx, &rvboxv1.GetClientRequest{ClientId: clientID})
+		if err != nil {
+			return err
+		}
+		item := response.GetClient()
+		fmt.Fprintf(output, "client %s connected=%t platform=%s running=%d queued=%d instance=%s\n", item.GetClientId(), item.GetConnected(), item.GetPlatform(), item.GetRunningCommands(), item.GetQueuedCommands(), item.GetClientInstanceId())
+		if item.GetPendingInstanceId() != "" {
+			fmt.Fprintf(output, "pending_instance=%s\n", item.GetPendingInstanceId())
+		}
+		return nil
+	}
+	response, err := client.GetCommand(ctx, &rvboxv1.GetCommandRequest{ClientId: clientID, IssueUuid: positionals[1]})
+	if err != nil {
+		return err
+	}
+	item := response.GetCommand()
+	fmt.Fprintf(output, "command %s client=%s lifecycle=%s revision=%d events=%d\n", item.GetIssueUuid(), item.GetTargetClientId(), item.GetLifecycle(), item.GetCommandRevision(), item.GetLastEventSeq())
+	if item.GetQueueExpiryTime() != nil {
+		fmt.Fprintf(output, "queue_expiry=%s\n", item.GetQueueExpiryTime().AsTime().UTC().Format(time.RFC3339Nano))
+	}
+	if item.GetTerminalTime() != nil {
+		fmt.Fprintf(output, "terminal_time=%s\n", item.GetTerminalTime().AsTime().UTC().Format(time.RFC3339Nano))
+	}
+	return nil
+}
+
+func runCommand(ctx context.Context, client rvboxv1.ControlClient, args []string, output, diagnostics io.Writer) error {
+	flags := flag.NewFlagSet("run", flag.ContinueOnError)
+	flags.SetOutput(io.Discard)
+	background := flags.Bool("background", false, "return after durable admission")
+	cwd := flags.String("cwd", "", "command working directory")
+	shell := flags.String("shell", "", "shell (sh, bash, cmd, powershell)")
+	requestID := flags.String("request-id", "", "canonical UUIDv7 used for idempotent admission")
+	queueTTL := flags.Duration("queue-ttl", -1, "queue TTL; zero means no expiry")
+	scriptPath := flags.String("script", "", "script file path")
+	envValues := repeatedFlag{}
+	profileValues := repeatedFlag{}
+	flags.Var(&envValues, "env", "environment override KEY=VALUE")
+	flags.Var(&profileValues, "profile", "execution profile")
+	if err := flags.Parse(args); err != nil {
+		return err
+	}
+	positionals := flags.Args()
+	if len(positionals) == 0 {
+		return errors.New("run requires CLIENT and COMMAND (or --script PATH CLIENT)")
+	}
+	if *requestID == "" {
+		generated, err := domain.NewUUIDv7()
+		if err != nil {
+			return err
+		}
+		*requestID = generated.String()
+	}
+	if _, err := domain.ParseUUIDv7(*requestID); err != nil {
+		return fmt.Errorf("--request-id: %w", err)
+	}
+	spec := &rvboxv1.ExecutionSpec{Cwd: *cwd}
+	spec.ShellType, _ = parseShell(*shell)
+	for _, value := range envValues {
+		parts := strings.SplitN(value, "=", 2)
+		if len(parts) != 2 || parts[0] == "" {
+			return fmt.Errorf("--env must be KEY=VALUE: %q", value)
+		}
+		if spec.EnvOverrides == nil {
+			spec.EnvOverrides = make(map[string]string)
+		}
+		spec.EnvOverrides[parts[0]] = parts[1]
+	}
+	for _, value := range profileValues {
+		profile, err := parseProfile(value)
+		if err != nil {
+			return err
+		}
+		spec.ExecutionProfiles = append(spec.ExecutionProfiles, profile)
+	}
+	var clientID string
+	if *scriptPath != "" {
+		if len(positionals) != 1 {
+			return errors.New("script form accepts exactly CLIENT")
+		}
+		clientID = positionals[0]
+		data, err := os.ReadFile(*scriptPath)
+		if err != nil {
+			return err
+		}
+		if len(data) > 10<<20 {
+			return errors.New("script exceeds the 10 MiB limit")
+		}
+		digest := sha256.Sum256(data)
+		spec.Source = &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: filepath.Base(*scriptPath), SizeBytes: uint64(len(data)), Sha256: digest[:]}}
+		request := &rvboxv1.RunCommandRequest{TargetClientId: clientID, Spec: spec, ScriptContent: data, RequestId: *requestID, QueueTtl: queueDuration(*queueTTL)}
+		return printRunResponse(ctx, client, request, output, *background)
+	}
+	if len(positionals) < 2 {
+		return errors.New("run requires CLIENT and COMMAND")
+	}
+	clientID = positionals[0]
+	spec.Source = &rvboxv1.ExecutionSpec_CommandText{CommandText: strings.Join(positionals[1:], " ")}
+	request := &rvboxv1.RunCommandRequest{TargetClientId: clientID, Spec: spec, RequestId: *requestID, QueueTtl: queueDuration(*queueTTL)}
+	_ = diagnostics
+	return printRunResponse(ctx, client, request, output, *background)
+}
+
+func printRunResponse(ctx context.Context, client rvboxv1.ControlClient, request *rvboxv1.RunCommandRequest, output io.Writer, _ bool) error {
+	response, err := client.RunCommand(ctx, request)
+	if err != nil {
+		return err
+	}
+	fmt.Fprintf(output, "%s %s\n", response.GetIssueUuid(), response.GetLifecycle())
+	return nil
+}
+
+func queueDuration(value time.Duration) *durationpb.Duration {
+	if value < 0 {
+		return nil
+	}
+	return durationpb.New(value)
+}
+
+func parseShell(value string) (rvboxv1.ShellType, error) {
+	switch strings.ToLower(value) {
+	case "", "default":
+		return rvboxv1.ShellType_SHELL_TYPE_UNSPECIFIED, nil
+	case "sh":
+		return rvboxv1.ShellType_SHELL_SH, nil
+	case "bash":
+		return rvboxv1.ShellType_SHELL_BASH, nil
+	case "cmd":
+		return rvboxv1.ShellType_SHELL_CMD, nil
+	case "powershell", "pwsh":
+		return rvboxv1.ShellType_SHELL_POWERSHELL, nil
+	default:
+		return rvboxv1.ShellType_SHELL_TYPE_UNSPECIFIED, fmt.Errorf("unknown shell %q", value)
+	}
+}
+
+func parseProfile(value string) (rvboxv1.ExecutionProfile, error) {
+	switch strings.ToLower(value) {
+	case "light":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_LIGHT, nil
+	case "cpu-medium":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_CPU_MEDIUM, nil
+	case "cpu-heavy":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_CPU_HEAVY, nil
+	case "mem-medium":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_MEM_MEDIUM, nil
+	case "mem-heavy":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_MEM_HEAVY, nil
+	case "disk-medium":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_DISK_MEDIUM, nil
+	case "disk-heavy":
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_DISK_HEAVY, nil
+	default:
+		return rvboxv1.ExecutionProfile_EXECUTION_PROFILE_UNSPECIFIED, fmt.Errorf("unknown profile %q", value)
+	}
+}
+
+type repeatedFlag []string
+
+func (flag *repeatedFlag) String() string { return strings.Join(*flag, ",") }
+func (flag *repeatedFlag) Set(value string) error {
+	*flag = append(*flag, value)
+	return nil
+}