diff --git a/sdks/go/test/integration/expansions.go b/sdks/go/test/integration/expansions.go index 2becce67333d..8565c1b3b153 100644 --- a/sdks/go/test/integration/expansions.go +++ b/sdks/go/test/integration/expansions.go @@ -16,13 +16,19 @@ package integration import ( + "context" "fmt" - "net" "strconv" "time" + jobpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1" "github.com/apache/beam/sdks/v2/go/test/integration/internal/jars" "github.com/apache/beam/sdks/v2/go/test/integration/internal/ports" + "google.golang.org/grpc" + "google.golang.org/grpc/backoff" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" ) // ExpansionServices is a struct used for getting addresses and starting expansion services, based @@ -91,47 +97,85 @@ func (es *ExpansionServices) GetAddr(label string) (string, error) { return "", fmt.Errorf("expansion service labeled \"%s\" not found: %w", label, err) } - // Start jar on open port. - port, err := ports.GetOpenTCP() - if err != nil { - return "", fmt.Errorf("cannot get open port for expansion service labeled \"%s\": %w", label, err) + // A different service may claim the port between selection and jar startup. + const startupAttempts = 3 + for attempt := 0; ; attempt++ { + port, err := ports.GetOpenTCP() + if err != nil { + return "", fmt.Errorf("cannot get open port for expansion service labeled \"%s\": %w", label, err) + } + portStr := strconv.Itoa(port) + proc, err := es.run(*ExpansionTimeout, jar, portStr) + if err != nil { + return "", fmt.Errorf("cannot run jar for expansion service labeled \"%s\": %w", label, err) + } + addr := "localhost:" + portStr + if es.testMode { + time.Sleep(es.waitTime) + } else { + exited := make(chan error, 1) + go func() { exited <- proc.Wait() }() + if err := waitForExpansionService(addr, exited); err != nil { + proc.Kill() + // Retry only startup collisions, never expansion or pipeline RPCs. + if status.Code(err) == codes.Unimplemented && attempt+1 < startupAttempts { + continue + } + return "", fmt.Errorf("expansion service labeled \"%s\" failed to start at %s: %w", label, addr, err) + } + } + es.procs = append(es.procs, proc) + es.addrs[label] = addr + return addr, nil } - portStr := strconv.Itoa(port) +} - // Run jar and cache its info. - proc, err := es.run(*ExpansionTimeout, jar, portStr) +func waitForExpansionService(addr string, exited <-chan error) error { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + conn, err := grpc.NewClient("passthrough:///"+addr, + grpc.WithTransportCredentials(insecure.NewCredentials()), + // The default reconnect backoff grows longer than the polling interval. + // Use a fixed delay so that the probe sees a late bind before the deadline. + grpc.WithConnectParams(grpc.ConnectParams{ + Backoff: backoff.Config{BaseDelay: 250 * time.Millisecond, Multiplier: 1, MaxDelay: 250 * time.Millisecond}, + MinConnectTimeout: 5 * time.Second, + })) if err != nil { - return "", fmt.Errorf("cannot run jar for expansion service labeled \"%s\": %w", label, err) + return err } - - addr := "localhost:" + portStr - - // Use different wait strategies for test mode vs production - if es.testMode { - // In test mode, use simple wait time for compatibility with mock processes - time.Sleep(es.waitTime) - } else { - // In production, wait for the jar to start with improved retry logic - maxRetries := 30 - retryDelay := time.Second - - for i := 0; i < maxRetries; i++ { - time.Sleep(retryDelay) - // Try to connect to the expansion service to verify it's ready - conn, err := net.DialTimeout("tcp", addr, 2*time.Second) - if err == nil { - conn.Close() - break - } - if i == maxRetries-1 { - return "", fmt.Errorf("expansion service labeled \"%s\" failed to start after %d retries: %w", label, maxRetries, err) + defer conn.Close() + client := jobpb.NewExpansionServiceClient(conn) + for { + probeCtx, probeCancel := context.WithTimeout(ctx, time.Second) + _, err := client.DiscoverSchemaTransform(probeCtx, &jobpb.DiscoverSchemaTransformRequest{}, grpc.WaitForReady(true)) + probeCancel() + if status.Code(err) == codes.Unimplemented { + // Older services only implement Expand. An empty request cannot expand + // a transform, but an ExpansionResponse (including its error field) + // proves that the service implements the expected RPC. Allow the + // remaining startup deadline for first-request initialization. + _, err = client.Expand(ctx, &jobpb.ExpansionRequest{}, grpc.WaitForReady(true)) + } + if err == nil { + select { + case err := <-exited: + return fmt.Errorf("jar exited before becoming ready: %v", err) + default: + return nil } } + if status.Code(err) != codes.Unavailable && status.Code(err) != codes.DeadlineExceeded { + return err + } + select { + case err := <-exited: + return fmt.Errorf("jar exited before becoming ready: %v", err) + case <-ctx.Done(): + return fmt.Errorf("waiting for expansion API: %w", ctx.Err()) + case <-time.After(time.Second): + } } - - es.procs = append(es.procs, proc) - es.addrs[label] = addr - return addr, nil } // Shutdown shuts down any jars started by the ExpansionServices struct and should get called if it diff --git a/sdks/go/test/integration/expansions_test.go b/sdks/go/test/integration/expansions_test.go index 6e47f5d303b0..1a42e90d0c54 100644 --- a/sdks/go/test/integration/expansions_test.go +++ b/sdks/go/test/integration/expansions_test.go @@ -16,28 +16,185 @@ package integration import ( + "context" "fmt" + "net" + "strings" + "sync/atomic" "testing" "time" + jobpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1" _ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow" _ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/flink" _ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/spark" "github.com/apache/beam/sdks/v2/go/test/integration/internal/jars" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" + "google.golang.org/grpc" ) type testProcess struct { killed bool jar string + done chan error } func (p *testProcess) Kill() error { p.killed = true + if p.done != nil { + select { + case <-p.done: + default: + close(p.done) + } + } return nil } +func (p *testProcess) Wait() error { + return <-p.done +} + +type readyExpansionService struct { + jobpb.UnimplementedExpansionServiceServer +} + +func (readyExpansionService) DiscoverSchemaTransform(context.Context, *jobpb.DiscoverSchemaTransformRequest) (*jobpb.DiscoverSchemaTransformResponse, error) { + return &jobpb.DiscoverSchemaTransformResponse{}, nil +} + +type expandOnlyService struct { + jobpb.UnimplementedExpansionServiceServer + delay time.Duration + calls atomic.Int32 +} + +func (s *expandOnlyService) Expand(ctx context.Context, _ *jobpb.ExpansionRequest) (*jobpb.ExpansionResponse, error) { + s.calls.Add(1) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(s.delay): + } + return &jobpb.ExpansionResponse{Error: "missing transform specification"}, nil +} + +func startExpansionListener(t *testing.T, port string, expansion jobpb.ExpansionServiceServer) string { + t.Helper() + listener, err := net.Listen("tcp", "localhost:"+port) + if err != nil { + // Another process can claim the selected port first. These tests cannot choose the port. + t.Skipf("cannot bind the selected port: %v", err) + } + server := grpc.NewServer() + if expansion != nil { + jobpb.RegisterExpansionServiceServer(server, expansion) + } + t.Cleanup(server.Stop) + go server.Serve(listener) + return listener.Addr().String() +} + +func TestExpansionServices_GetAddr_RejectsWrongService(t *testing.T) { + var procs []*testProcess + es := &ExpansionServices{ + addrs: map[string]string{}, + jars: map[string]string{"test": "test.jar"}, + run: func(_ time.Duration, _ string, args ...string) (jars.Process, error) { + // Another service claims the selected port before the jar can bind it. + startExpansionListener(t, args[0], nil) + proc := &testProcess{done: make(chan error)} + procs = append(procs, proc) + return proc, nil + }, + } + t.Cleanup(es.Shutdown) + if addr, err := es.GetAddr("test"); err == nil { + t.Fatalf("GetAddr() = %q, nil; want an error for a different gRPC service", addr) + } + if len(es.addrs) != 0 || len(es.procs) != 0 { + t.Fatalf("GetAddr() cached an unsuccessful startup: addrs=%v, procs=%v", es.addrs, es.procs) + } + for _, proc := range procs { + if !proc.killed { + t.Error("GetAddr() did not kill the unsuccessful jar process") + } + } + if len(procs) != 3 { + t.Errorf("GetAddr() started %d jars, want 3 bounded startup attempts", len(procs)) + } +} + +func TestExpansionServices_GetAddr_RetriesWrongService(t *testing.T) { + var procs []*testProcess + var ports []string + es := &ExpansionServices{ + addrs: map[string]string{}, + jars: map[string]string{"test": "test.jar"}, + run: func(_ time.Duration, _ string, args ...string) (jars.Process, error) { + ports = append(ports, args[0]) + var service jobpb.ExpansionServiceServer + if len(procs) > 0 { + service = readyExpansionService{} + } + startExpansionListener(t, args[0], service) + proc := &testProcess{done: make(chan error, 1)} + if len(procs) == 0 { + // The jar cannot bind the port held by the other gRPC service. + proc.done <- nil + close(proc.done) + } + procs = append(procs, proc) + return proc, nil + }, + } + t.Cleanup(es.Shutdown) + addr, err := es.GetAddr("test") + if err != nil { + t.Fatal(err) + } + if len(ports) != 2 || ports[0] == ports[1] { + t.Fatalf("GetAddr() selected ports %v, want a fresh port after the collision", ports) + } + if addr != "localhost:"+ports[1] || !procs[0].killed || procs[1].killed { + t.Errorf("GetAddr() = %q, killed processes = [%v, %v]", addr, procs[0].killed, procs[1].killed) + } + if got, err := es.GetAddr("test"); err != nil || got != addr || len(ports) != 2 { + t.Errorf("cached GetAddr() = %q, %v; ports = %v", got, err, ports) + } +} + +func TestWaitForExpansionService_RejectsExitedJar(t *testing.T) { + // A compatible listener still cannot prove that our child is running. + addr := startExpansionListener(t, "0", readyExpansionService{}) + exited := make(chan error, 1) + exited <- fmt.Errorf("jar failed to start") + if err := waitForExpansionService(addr, exited); err == nil || !strings.Contains(err.Error(), "jar failed to start") { + t.Errorf("waitForExpansionService() = %v; want the child startup error", err) + } +} + +func TestExpansionServices_GetAddr_ExpandOnlyService(t *testing.T) { + // The first expansion can initialize classes after the server binds. + service := &expandOnlyService{delay: 1200 * time.Millisecond} + es := &ExpansionServices{ + addrs: map[string]string{}, + jars: map[string]string{"test": "test.jar"}, + run: func(_ time.Duration, _ string, args ...string) (jars.Process, error) { + startExpansionListener(t, args[0], service) + return &testProcess{done: make(chan error)}, nil + }, + } + t.Cleanup(es.Shutdown) + if addr, err := es.GetAddr("test"); err != nil || addr == "" { + t.Errorf("GetAddr() = %q, %v; want the Expand-only service address", addr, err) + } + if got := service.calls.Load(); got != 1 { + t.Errorf("Expand called %d times, want a single startup probe", got) + } +} + func failRun(_ time.Duration, _ string, _ ...string) (jars.Process, error) { return nil, fmt.Errorf("unexpectedly running a jar, failing") } diff --git a/sdks/go/test/integration/internal/jars/proc.go b/sdks/go/test/integration/internal/jars/proc.go index fa12beeead9a..0194f75fd278 100644 --- a/sdks/go/test/integration/internal/jars/proc.go +++ b/sdks/go/test/integration/internal/jars/proc.go @@ -15,8 +15,19 @@ package jars -// Process is an interface to allow wrapping os.Process with alternate behavior +import "os/exec" + +// Process is an interface to allow wrapping exec.Cmd with alternate behavior // depending on OS. type Process interface { Kill() error + Wait() error +} + +type commandProcess struct { + *exec.Cmd +} + +func (p *commandProcess) Kill() error { + return p.Process.Kill() } diff --git a/sdks/go/test/integration/internal/jars/proc_unix.go b/sdks/go/test/integration/internal/jars/proc_unix.go index c9f2d9d05bea..a835b4589cf4 100644 --- a/sdks/go/test/integration/internal/jars/proc_unix.go +++ b/sdks/go/test/integration/internal/jars/proc_unix.go @@ -19,15 +19,12 @@ package jars -import ( - "golang.org/x/sys/unix" - "os" -) +import "golang.org/x/sys/unix" -// UnixProcess wraps os.Process and changes the Kill function to perform a more +// UnixProcess wraps exec.Cmd and changes the Kill function to perform a more // graceful shutdown, mainly for compatibility with the timeout command. type UnixProcess struct { - proc *os.Process // The os.Process for the running jar. + commandProcess } // Kill gracefully shuts down the process. It is recommended to use this @@ -35,5 +32,5 @@ type UnixProcess struct { func (p *UnixProcess) Kill() error { // Avoid using SIGKILL. If the jar is wrapped in the timeout command // SIGKILL will kill the timeout and leave the jar running. - return p.proc.Signal(unix.SIGTERM) + return p.Process.Signal(unix.SIGTERM) } diff --git a/sdks/go/test/integration/internal/jars/run_nonunix.go b/sdks/go/test/integration/internal/jars/run_nonunix.go index 2d392e7cfa2d..d306dfb2f6a3 100644 --- a/sdks/go/test/integration/internal/jars/run_nonunix.go +++ b/sdks/go/test/integration/internal/jars/run_nonunix.go @@ -40,7 +40,7 @@ func getTimeoutRunner() runCallback { } } -// run simply starts up a jar and returns the cmd.Process. +// run simply starts up a jar and returns its process. func run(jar string, args ...string) (Process, error) { var cmdArr []string cmdArr = append(cmdArr, "java", "-jar", jar) @@ -51,5 +51,5 @@ func run(jar string, args ...string) (Process, error) { if err != nil { return nil, err } - return cmd.Process, nil + return &commandProcess{cmd}, nil } diff --git a/sdks/go/test/integration/internal/jars/run_unix.go b/sdks/go/test/integration/internal/jars/run_unix.go index 2a401b888ec8..986a1a99ec4b 100644 --- a/sdks/go/test/integration/internal/jars/run_unix.go +++ b/sdks/go/test/integration/internal/jars/run_unix.go @@ -61,5 +61,5 @@ func run(dur time.Duration, jar string, args ...string) (Process, error) { if err != nil { return nil, err } - return &UnixProcess{proc: cmd.Process}, nil + return &UnixProcess{commandProcess{cmd}}, nil }