[Go] Verify expansion service readiness before use

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

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")
 }

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()
 }

sdks/go/test/integration/internal/jars/proc_unix.go

@@ -19,21 +19,18 @@
 
 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
 // instead of directly killing the process.
 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)
 }

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
 }

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
 }