You are an expert in writing code in many different languages. Your goal is to perform code completion for the following code, keeping in mind the rest of the code and the file meta data.
    Metadata details:
    filename: stage.go
    path of file: ./repos/beam/sdks/go/pkg/beam/runners/prism/internal
    
    the code of the file until where you have to start completion: // Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package internal

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"time"

	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine"
	"github.com/apache/beam/sdks/v2/go/pkg/be
    
func (s *stage) Execute(ctx context.Context, j *jobservices.Job, wk *worker.W, comps *pipepb.Components, em *engine.ElementManager, rb engine.RunBundle) error {
	slog.Debug("Execute: starting bundle", "bundle", rb)

	var b *worker.B
	initialState := em.StateForBundle(rb)
	var dataReady <-chan struct{}
	switch s.envID {
	case "": // Runner Transforms
		if len(s.transforms) != 1 {
			panic(fmt.Sprintf("unexpected number of runner transforms, want 1: %+v", s))
		}
		tid := s.transforms[0]
		// Runner transforms are processed immeadiately.
		b = s.exe.ExecuteTransform(s.ID, tid, comps.GetTransforms()[tid], comps, rb.Watermark, em.InputForBundle(rb, s.inputInfo))
		b.InstID = rb.BundleID
		slog.Debug("Execute: runner transform", "bundle", rb, slog.String("tid", tid))

		// Do some accounting for the fake bundle.
		b.Resp = make(chan *fnpb.ProcessBundleResponse, 1)
		close(b.Resp) // To avoid blocking downstream, since we don't send on this.
		closed := make(chan struct{})
		close(closed)
		dataReady = closed
	case wk.Env:
		input, estimatedElements := em.DataAndTimerInputForBundle(rb, s.inputInfo)
		b = &worker.B{
			PBDID:  s.ID,
			InstID: rb.BundleID,

			InputTransformID: s.inputTransformID,

			Input:                  input,
			EstimatedInputElements: estimatedElements,

			OutputData: initialState,
			HasTimers:  s.hasTimers,

			SinkToPCollection: s.SinkToPCollection,
			OutputCount:       len(s.outputs),
		}
		b.Init()

		s.prepareSides(b, rb.Watermark)

		slog.Debug("Execute: processing", "bundle", rb)
		defer b.Cleanup(wk)
		dataReady = b.ProcessOn(ctx, wk)
	default:
		err := fmt.Errorf("unknown environment[%v]", s.envID)
		slog.Error("Execute", "error", err)
		panic(err)
	}

	// Progress + split loop.
	previousIndex := int64(-2)
	previousTotalCount := int64(-2) // Total count of all pcollection elements.

	unsplit := true
	progTick := time.NewTicker(100 * time.Millisecond)
	defer progTick.Stop()
	var dataFinished, bundleFinished bool
	// If we have no data outputs, we still need to have progress & splits
	// while waiting for bundle completion.
	if b.OutputCount == 0 {
		dataFinished = true
	}
	var resp *fnpb.ProcessBundleResponse
progress:
	for {
		select {
		case <-ctx.Done():
			return context.Cause(ctx)
		case resp = <-b.Resp:
			bundleFinished = true
			if b.BundleErr != nil {
				return b.BundleErr
			}
			if dataFinished && bundleFinished {
				break progress // exit progress loop on close.
			}
		case <-dataReady:
			dataFinished = true
			if dataFinished && bundleFinished {
				break progress // exit progress loop on close.
			}
		case <-progTick.C:
			resp, err := b.Progress(ctx, wk)
			if err != nil {
				slog.Debug("SDK Error from progress, aborting progress", "bundle", rb, "error", err.Error())
				break progress
			}
			index, unknownIDs := j.ContributeTentativeMetrics(resp)
			if len(unknownIDs) > 0 {
				md := wk.MonitoringMetadata(ctx, unknownIDs)
				j.AddMetricShortIDs(md)
			}
			slog.Debug("progress report", "bundle", rb, "index", index, "prevIndex", previousIndex)

			// Check if there has been any measurable progress by the input, or all output pcollections since last report.
			slow := previousIndex == index["index"] && previousTotalCount == index["totalCount"]
			if slow && unsplit {
				slog.Debug("splitting report", "bundle", rb, "index", index)
				sr, err := b.Split(ctx, wk, 0.5 /* fraction of remainder */, nil /* allowed splits */)
				if err != nil {
					slog.Warn("SDK Error from split, aborting splits", "bundle", rb, "error", err.Error())
					break progress
				}
				if sr.GetChannelSplits() == nil {
					slog.Debug("SDK returned no splits", "bundle", rb)
					unsplit = false
					continue progress
				}
				// TODO sort out rescheduling primary Roots on bundle failure.
				var residualData [][]byte
				for _, rr := range sr.GetResidualRoots() {
					ba := rr.GetApplication()
					residualData = append(residualData, ba.GetElement())
					if len(ba.GetElement()) == 0 {
						slog.LogAttrs(context.TODO(), slog.LevelError, "returned empty residual application", slog.Any("bundle", rb))
						panic("sdk returned empty residual application")
					}
					// TODO what happens to output watermarks on splits?
				}
				if len(sr.GetChannelSplits()) != 1 {
					slog.Warn("received non-single channel split", "bundle", rb)
				}
				cs := sr.GetChannelSplits()[0]
				fr := cs.GetFirstResidualElement()
				// The first residual can be after the end of data, so filter out those cases.
				if b.EstimatedInputElements >= int(fr) {
					b.EstimatedInputElements = int(fr) // Update the estimate for the next split.
					// Residuals are returned right away for rescheduling.
					em.ReturnResiduals(rb, int(fr), s.inputInfo, residualData)
				}
			} else {
				previousIndex = index["index"]
				previousTotalCount = index["totalCount"]
			}
		}
	}
	// Tentative Data is ready, commit it to the main datastore.
	slog.Debug("Execute: committing data", "bundle", rb, slog.Any("outputsWithData", maps.Keys(b.OutputData.Raw)), slog.Any("outputs", maps.Keys(s.OutputsToCoders)))

	// Tally metrics immeadiately so they're available before
	// pipeline termination.
	unknownIDs := j.ContributeFinalMetrics(resp)
	if len(unknownIDs) > 0 {
		md := wk.MonitoringMetadata(ctx, unknownIDs)
		j.AddMetricShortIDs(md)
	}
	var residualData [][]byte
	var minOutputWatermark map[string]mtime.Time
	for _, rr := range resp.GetResidualRoots() {
		ba := rr.GetApplication()
		residualData = append(residualData, ba.GetElement())
		if len(ba.GetElement()) == 0 {
			slog.LogAttrs(context.TODO(), slog.LevelError, "returned empty residual application", slog.Any("bundle", rb))
			panic("sdk returned empty residual application")
		}
		for col, wm := range ba.GetOutputWatermarks() {
			if minOutputWatermark == nil {
				minOutputWatermark = map[string]mtime.Time{}
			}
			cur, ok := minOutputWatermark[col]
			if !ok {
				cur = mtime.MaxTimestamp
			}
			minOutputWatermark[col] = mtime.Min(mtime.FromTime(wm.AsTime()), cur)
		}
	}
	if l := len(residualData); l > 0 {
		slog.Debug("returned empty residual application", "bundle", rb, slog.Int("numResiduals", l), slog.String("pcollection", s.primaryInput))
	}
	em.PersistBundle(rb, s.OutputsToCoders, b.OutputData, s.inputInfo, residualData, minOutputWatermark)
	b.OutputData = engine.TentativeData{} // Clear the data.
	return nil
}