refactor(client): make ProxySingleMachineContext and ProxyMachinesContext package functions as well

This commit is contained in:
Pasha Sviderski
2026-08-28 20:37:50 +10:00
parent 43bf2baaf7
commit 00d68d9465
8 changed files with 114 additions and 16 deletions
+1 -1
View File
@@ -187,7 +187,7 @@ func printContainers(containers []containerInfo) error {
}
func collectContainers(ctx context.Context, cli *client.Client) ([]containerInfo, error) {
listCtx := cli.ProxyMachinesContext(ctx, nil)
listCtx := client.ProxyMachinesContext(ctx, nil)
// List all service containers across all machines in the cluster.
machineContainers, err := cli.Docker.ListServiceContainers(
+23 -4
View File
@@ -78,8 +78,8 @@ func (cli *Client) progressOut() *streams.Out {
// ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines.
// If namesOrIDs is nil or empty, all machines are included.
// This triggers One2Many proxying, which always injects metadata into the response.
func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
md := metadata.New(nil)
func ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
md := outgoingMetadataWithoutProxyTargets(ctx)
if len(namesOrIDs) == 0 {
md.Append("machines", "*")
} else {
@@ -92,7 +92,26 @@ func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string
// ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine.
// This triggers One2One proxying, which does NOT inject metadata into the response.
// Use this for requests that expect a single response message without metadata wrapper.
func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
md := metadata.Pairs("machine", nameOrID)
func ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
md := outgoingMetadataWithoutProxyTargets(ctx)
md.Set("machine", nameOrID)
return metadata.NewOutgoingContext(ctx, md)
}
func outgoingMetadataWithoutProxyTargets(ctx context.Context) metadata.MD {
md, _ := metadata.FromOutgoingContext(ctx)
md = md.Copy()
md.Delete("machine")
md.Delete("machines")
return md
}
// ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines.
func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
return ProxyMachinesContext(ctx, namesOrIDs)
}
// ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine.
func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
return ProxySingleMachineContext(ctx, nameOrID)
}
+79
View File
@@ -0,0 +1,79 @@
package client
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/metadata"
)
func TestProxySingleMachineContext(t *testing.T) {
original := metadata.Pairs(
"authorization", "token",
"machine", "old-machine",
"machines", "old-machine-a",
"machines", "old-machine-b",
)
ctx := metadata.NewOutgoingContext(context.Background(), original)
proxyCtx := ProxySingleMachineContext(ctx, "new-machine")
md, ok := metadata.FromOutgoingContext(proxyCtx)
require.True(t, ok)
require.Equal(t, metadata.Pairs(
"authorization", "token",
"machine", "new-machine",
), md)
require.Equal(t, metadata.Pairs(
"authorization", "token",
"machine", "old-machine",
"machines", "old-machine-a",
"machines", "old-machine-b",
), original)
}
func TestProxyMachinesContext(t *testing.T) {
tests := []struct {
name string
machines []string
want []string
}{
{
name: "specified machines",
machines: []string{"machine-a", "machine-b"},
want: []string{"machine-a", "machine-b"},
},
{
name: "all machines",
want: []string{"*"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
original := metadata.Pairs(
"authorization", "token",
"machine", "old-machine",
"machines", "old-machine-a",
"machines", "old-machine-b",
)
ctx := metadata.NewOutgoingContext(context.Background(), original)
proxyCtx := ProxyMachinesContext(ctx, tt.machines)
md, ok := metadata.FromOutgoingContext(proxyCtx)
require.True(t, ok)
require.Equal(t, metadata.MD{
"authorization": {"token"},
"machines": tt.want,
}, md)
require.Equal(t, metadata.Pairs(
"authorization", "token",
"machine", "old-machine",
"machines", "old-machine-a",
"machines", "old-machine-b",
), original)
})
}
}
+4 -4
View File
@@ -74,7 +74,7 @@ func (cli *Client) createServiceContainerWithPull(
resp.Name = containerName
// Proxy Docker gRPC requests to the selected machine.
ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
pw := progress.ContextWriter(ctx)
eventID := cliprogress.NewContainerEventID(ctx, containerName, machine.Machine.Name)
@@ -277,7 +277,7 @@ func (cli *Client) resolveContainerOperation(
eventID := cliprogress.ContainerEventID(ctx, ctr.Container.ServiceSpec.Name, ctr.Container.ID, ctr.MachineName)
return containerOperationContext{
ctx: cli.ProxySingleMachineContext(ctx, ctr.MachineID),
ctx: ProxySingleMachineContext(ctx, ctr.MachineID),
containerID: ctr.Container.ID,
eventID: eventID,
}, nil
@@ -375,7 +375,7 @@ func (cli *Client) ExecContainer(
}
// Proxy Docker gRPC requests to the machine hosting the container
ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
// Execute the command in the container
exitCode, err := cli.Docker.ExecContainer(ctx, machinedocker.ExecConfig{
@@ -452,7 +452,7 @@ func (cli *Client) WaitContainerHealthy(
}
// For containers with a health check, wait until Docker reports healthy or unhealthy.
mctx := cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
mctx := ProxySingleMachineContext(ctx, machine.Machine.Id)
mctx, cancel := context.WithTimeout(mctx, healthcheckTimeout(mc.Container.Config.Healthcheck))
defer cancel()
ticker := time.NewTicker(1 * time.Second)
+1 -1
View File
@@ -59,7 +59,7 @@ func (cli *Client) InspectRemoteImage(ctx context.Context, id string) ([]api.Mac
// it lists images on all machines.
func (cli *Client) ListImages(ctx context.Context, filter api.ImageFilter) ([]api.MachineImages, error) {
// Broadcast the image list request to the specified machines or all machines if none specified.
listCtx := cli.ProxyMachinesContext(ctx, filter.Machines)
listCtx := ProxyMachinesContext(ctx, filter.Machines)
opts := image.ListOptions{Manifests: true}
if filter.Name != "" {
+2 -2
View File
@@ -102,7 +102,7 @@ func (cli *Client) ServiceLogs(
func (cli *Client) ContainerLogs(
ctx context.Context, machineNameOrID string, containerID string, opts api.ServiceLogsOptions,
) (<-chan api.LogEntry, error) {
proxyCtx := cli.ProxySingleMachineContext(ctx, machineNameOrID)
proxyCtx := ProxySingleMachineContext(ctx, machineNameOrID)
req := &pb.LogsRequest{
Id: containerID,
@@ -198,7 +198,7 @@ func (cli *Client) MachineLogs(
func (cli *Client) systemServiceLogs(
ctx context.Context, machineID, service string, opts api.ServiceLogsOptions,
) (<-chan api.LogEntry, error) {
proxyCtx := cli.ProxySingleMachineContext(ctx, machineID)
proxyCtx := ProxySingleMachineContext(ctx, machineID)
req := &pb.LogsRequest{
Id: service,
+1 -1
View File
@@ -78,7 +78,7 @@ func (cli *Client) ListMachines(ctx context.Context, filter *api.MachineFilter)
func (cli *Client) UpdateMachine(
ctx context.Context, nameOrID string, req *pb.UpdateMachineRequest,
) (*pb.MachineInfo, error) {
ctx = cli.ProxySingleMachineContext(ctx, nameOrID)
ctx = ProxySingleMachineContext(ctx, nameOrID)
resp, err := cli.MachineClient.UpdateMachine(ctx, req)
if err != nil {
if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
+3 -3
View File
@@ -27,7 +27,7 @@ func (cli *Client) CreateVolume(
return resp, fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err)
}
// Proxy Docker gRPC requests to the selected machine.
ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
pw := progress.ContextWriter(ctx)
eventID := cliprogress.VolumeEventID(opts.Name, machine.Machine.Name)
@@ -56,7 +56,7 @@ func (cli *Client) ListVolumes(ctx context.Context, filter *api.VolumeFilter) ([
proxyMachines = filter.Machines
}
listCtx := cli.ProxyMachinesContext(ctx, proxyMachines)
listCtx := ProxyMachinesContext(ctx, proxyMachines)
machineVolumes, err := cli.Docker.ListVolumes(listCtx, volume.ListOptions{})
if err != nil {
return nil, err
@@ -107,7 +107,7 @@ func (cli *Client) RemoveVolume(ctx context.Context, machineNameOrID, volumeName
return fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err)
}
// Proxy Docker gRPC requests to the selected machine.
ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
pw := progress.ContextWriter(ctx)
eventID := cliprogress.VolumeEventID(volumeName, machine.Machine.Name)