[Prism] Fix resource (docker container) leak when shutting down (#35991)

* Fix prism resource leak when shutting down.

* Move wg into MultiplexW to avoid changing MakeWorker api

* Add timeout while waiting for WaitGroup to finish.

* Fix a worker_test.

* Make PendingDone a job method and use atomic internally.

* Fix comments for exported methods.

* Change MakkeWorker to take job id as the first argument. Checking if id is in wg before calling wait.

* Fix test

Author: shunping

sdks/go/pkg/beam/core/runtime/harness/harness.go

@@ -704,7 +704,7 @@ func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRe
 func parseTimeoutDurationFlag(ctx context.Context, elementProcessingTimeout string) time.Duration {
 	userSpecifiedTimeout, err := time.ParseDuration(elementProcessingTimeout)
 	if err != nil {
-		log.Errorf(ctx, "Failed to parse element_processing_timeout: %v, there will be no timeout for processing an element in a PTransform operation", err)
+		log.Warnf(ctx, "Failed to parse element_processing_timeout: %v, there will be no timeout for processing an element in a PTransform operation", err)
 		return 0 * time.Minute
 	}
 	return userSpecifiedTimeout

sdks/go/pkg/beam/runners/prism/internal/environments.go

@@ -259,7 +259,7 @@ func dockerEnvironment(ctx context.Context, logger *slog.Logger, dp *pipepb.Dock
 				defer rc.Close()
 				var buf bytes.Buffer
 				stdcopy.StdCopy(&buf, &buf, rc)
-				logger.Info("container being killed", slog.Any("cause", context.Cause(ctx)), slog.Any("containerLog", buf))
+				logger.Info("container being killed", slog.Any("cause", context.Cause(ctx)), slog.String("containerLog", buf.String()))
 			}
 			// Can't use command context, since it's already canceled here.
 			if err := cli.ContainerKill(bgctx, containerID, ""); err != nil {
@@ -288,6 +288,8 @@ func dockerEnvironment(ctx context.Context, logger *slog.Logger, dp *pipepb.Dock
 }
 
 func processEnvironment(ctx context.Context, logger *slog.Logger, pp *pipepb.ProcessPayload, wk *worker.W) {
+	defer wk.Stop()
+
 	cmd := exec.CommandContext(ctx, pp.GetCommand(), "--id='"+wk.ID+"'", "--provision_endpoint="+wk.Endpoint())
 	logger.Debug("starting process", "cmd", cmd.String())
 

sdks/go/pkg/beam/runners/prism/internal/execute.go

@@ -76,6 +76,7 @@ func RunPipeline(j *jobservices.Job) {
 	// any related job resources.
 	defer func() {
 		j.CancelFn(fmt.Errorf("runPipeline returned, cleaning up"))
+		j.WaitForCleanUp()
 	}()
 
 	j.SendMsg("running " + j.String())
@@ -95,7 +96,7 @@ func RunPipeline(j *jobservices.Job) {
 	j.SendMsg("pipeline completed " + j.String())
 
 	j.SendMsg("terminating " + j.String())
-	j.Done()
+	j.PendingDone()
 }
 
 type transformExecuter interface {

sdks/go/pkg/beam/runners/prism/internal/jobservices/job.go

@@ -94,6 +94,8 @@ type Job struct {
 	// Logger for this job.
 	Logger *slog.Logger
 
+	pendingDone atomic.Bool // indicate the job is done but waiting for clean-up
+
 	metrics metricsStore
 	mw      *worker.MultiplexW
 }
@@ -194,6 +196,20 @@ func (j *Job) Canceled() {
 	j.sendState(jobpb.JobState_CANCELLED)
 }
 
+// PendingDone indicates that the job is completed and is waiting for clean-up.
+func (j *Job) PendingDone() {
+	j.pendingDone.Store(true)
+}
+
+// WaitForCleanUp waits until all environments relevant to the job are cleaned up.
+func (j *Job) WaitForCleanUp() {
+	j.mw.WaitForCleanUp(j.String())
+	if j.pendingDone.Load() {
+		// If there is a pending done, only mark it as done after clean-up
+		j.Done()
+	}
+}
+
 // Failed indicates that the job completed unsuccessfully.
 func (j *Job) Failed(err error) {
 	slog.Error("job failed", slog.Any("job", j), slog.Any("error", err))
@@ -204,7 +220,7 @@ func (j *Job) Failed(err error) {
 
 // MakeWorker instantiates a worker.W populating environment and pipeline data from the Job.
 func (j *Job) MakeWorker(env string) *worker.W {
-	wk := j.mw.MakeWorker(j.String()+"_"+env, env)
+	wk := j.mw.MakeWorker(j.String(), env)
 	wk.EnvPb = j.Pipeline.GetComponents().GetEnvironments()[env]
 	wk.PipelineOptions = j.PipelineOptions()
 	wk.JobKey = j.JobKey()

sdks/go/pkg/beam/runners/prism/internal/worker/worker.go

@@ -29,6 +29,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
@@ -76,6 +77,7 @@ type W struct {
 	mu                 sync.Mutex
 	activeInstructions map[string]controlResponder              // Active instructions keyed by InstructionID
 	Descriptors        map[string]*fnpb.ProcessBundleDescriptor // Stages keyed by PBDID
+	wg                 *sync.WaitGroup
 }
 
 type controlResponder interface {
@@ -142,6 +144,7 @@ func (wk *W) shutdown() {
 func (wk *W) Stop() {
 	wk.shutdown()
 	wk.parentPool.delete(wk)
+	wk.wg.Done()
 	slog.Debug("stopped", "worker", wk)
 }
 
@@ -709,6 +712,7 @@ type MultiplexW struct {
 	endpoint string
 	logger   *slog.Logger
 	pool     map[string]*W
+	wg       map[string]*sync.WaitGroup
 }
 
 // NewMultiplexW instantiates a new FnAPI server for multiplexing FnAPI requests to a W.
@@ -718,6 +722,7 @@ func NewMultiplexW(lis net.Listener, g *grpc.Server, logger *slog.Logger) *Multi
 		endpoint: "localhost:" + p,
 		logger:   logger,
 		pool:     make(map[string]*W),
+		wg:       make(map[string]*sync.WaitGroup),
 	}
 
 	fnpb.RegisterBeamFnControlServer(g, mw)
@@ -735,8 +740,12 @@ func NewMultiplexW(lis net.Listener, g *grpc.Server, logger *slog.Logger) *Multi
 func (mw *MultiplexW) MakeWorker(id, env string) *W {
 	mw.mu.Lock()
 	defer mw.mu.Unlock()
+	workerId := id + "_" + env
+	if _, ok := mw.wg[id]; !ok {
+		mw.wg[id] = &sync.WaitGroup{}
+	}
 	w := &W{
-		ID:  id,
+		ID:  workerId,
 		Env: env,
 
 		InstReqs:    make(chan *fnpb.InstructionRequest, 10),
@@ -746,8 +755,11 @@ func (mw *MultiplexW) MakeWorker(id, env string) *W {
 		activeInstructions: make(map[string]controlResponder),
 		Descriptors:        make(map[string]*fnpb.ProcessBundleDescriptor),
 		parentPool:         mw,
+		wg:                 mw.wg[id],
 	}
-	mw.pool[id] = w
+	mw.pool[workerId] = w
+
+	mw.wg[id].Add(1)
 	return w
 }
 
@@ -809,6 +821,32 @@ func (mw *MultiplexW) delete(w *W) {
 	delete(mw.pool, w.ID)
 }
 
+// WaitForCleanUp waits until all resources relevant to the job are cleaned up.
+func (mw *MultiplexW) WaitForCleanUp(id string) {
+	mw.mu.Lock()
+	wg := mw.wg[id]
+	mw.mu.Unlock()
+	if wg == nil {
+		return
+	}
+
+	const cleanUpTimeout = 60 * time.Second
+	c := make(chan struct{})
+	go func() {
+		defer close(c)
+		wg.Wait()
+	}()
+
+	select {
+	case <-c: // Waitgroup finishes successfully
+		slog.Debug("Finished cleaning up job " + id)
+		return
+	case <-time.After(cleanUpTimeout): // Timeout
+		slog.Warn("Timeout when cleaning up job " + id)
+		return
+	}
+}
+
 func handleUnary[Request any, Response any, Method func(*W, context.Context, *Request) (*Response, error)](mw *MultiplexW, ctx context.Context, req *Request, m Method) (*Response, error) {
 	w, err := mw.workerFromMetadataCtx(ctx)
 	if err != nil {

sdks/go/pkg/beam/runners/prism/internal/worker/worker_test.go

@@ -44,7 +44,7 @@ func TestMultiplexW_MakeWorker(t *testing.T) {
 	if w.parentPool == nil {
 		t.Errorf("MakeWorker instantiated W with a nil reference to MultiplexW")
 	}
-	if got, want := w.ID, "test"; got != want {
+	if got, want := w.ID, "test_testEnv"; got != want {
 		t.Errorf("MakeWorker(%q) = %v, want %v", want, got, want)
 	}
 	got, ok := w.parentPool.pool[w.ID]
@@ -77,8 +77,8 @@ func TestMultiplexW_workerFromMetadataCtx(t *testing.T) {
 		},
 		{
 			name: "matched worker_id",
-			ctx:  metadata.NewIncomingContext(context.Background(), metadata.Pairs("worker_id", "test")),
-			want: &W{ID: "test"},
+			ctx:  metadata.NewIncomingContext(context.Background(), metadata.Pairs("worker_id", "test_testEnv")),
+			want: &W{ID: "test_testEnv"},
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
@@ -525,6 +525,7 @@ func TestWorker_State_MultimapSideInput(t *testing.T) {
 func newWorker() *W {
 	mw := &MultiplexW{
 		pool: map[string]*W{},
+		wg:   map[string]*sync.WaitGroup{},
 	}
 	return mw.MakeWorker("test", "testEnv")
 }