Add experimental worker.GetSlotUsage

CHANGELOG.md

@@ -47,6 +47,7 @@ to docs, or any other relevant information.
   standalone activities they start, which use it for payloads returned to the Nexus caller: results,
   failures, and standalone activity heartbeat details. When completing or heartbeating such an
   activity from a client, pass `activity.Info.NexusSerializationContext` in the options.
+- Added experimental `worker.GetSlotUsage`, which returns a snapshot of a worker's task slot usage.
 
 ### Changed
 

internal/internal_worker.go

@@ -1346,8 +1346,12 @@ type AggregatedWorker struct {
 	activityWorker *activityWorker
 	sessionWorker  *sessionWorker
 	nexusWorker    *nexusWorker
-	logger         log.Logger
-	registry       *registry
+	// The Nexus worker's slot supplier. Unlike the other workers, nexusWorker is created during
+	// start rather than at construction, so SlotUsage reads this atomic instead of nexusWorker to
+	// avoid a data race when called concurrently with start.
+	nexusSlotSupplier atomic.Pointer[trackingSlotSupplier]
+	logger            log.Logger
+	registry          *registry
 	// Stores a boolean indicating whether the worker has already been started.
 	started      atomic.Bool
 	shuttingDown atomic.Bool
@@ -1619,6 +1623,7 @@ func (aw *AggregatedWorker) start() error {
 		if err != nil {
 			return fmt.Errorf("failed to create a nexus worker: %w", err)
 		}
+		aw.nexusSlotSupplier.Store(aw.nexusWorker.worker.slotSupplier)
 		if err := aw.nexusWorker.Start(); err != nil {
 			return fmt.Errorf("failed to start a nexus worker: %w", err)
 		}

internal/slot_usage.go

@@ -0,0 +1,70 @@
+package internal
+
+import "time"
+
+// SlotUsage is a point-in-time snapshot of a worker's task slot usage, broken down by task kind.
+// Each task kind is read independently, so the snapshot is not atomic across task kinds.
+//
+// NOTE: Experimental
+//
+// Exposed as: [go.temporal.io/sdk/worker.SlotUsage]
+type SlotUsage struct {
+	// Workflow is the usage of workflow task slots.
+	Workflow SlotKindUsage
+	// Activity is the usage of activity task slots. It includes session activity and session
+	// creation tasks, if sessions are enabled. A session creation task holds its slot for the life
+	// of the session.
+	Activity SlotKindUsage
+	// LocalActivity is the usage of local activity slots.
+	LocalActivity SlotKindUsage
+	// Nexus is the usage of Nexus task slots.
+	Nexus SlotKindUsage
+}
+
+// SlotKindUsage is a snapshot of the task slot usage for one task kind.
+//
+// NOTE: Experimental
+//
+// Exposed as: [go.temporal.io/sdk/worker.SlotKindUsage]
+type SlotKindUsage struct {
+	// UsedSlots is the number of tasks currently being processed. Slots reserved by pollers that
+	// have not received a task are not counted.
+	UsedSlots int
+	// UsedSlotsLastChanged is when a task last started or finished, or the zero time if no task has
+	// started since the worker was created. A task that starts and finishes between two snapshots
+	// still changes it.
+	UsedSlotsLastChanged time.Time
+}
+
+func (u SlotKindUsage) add(other SlotKindUsage) SlotKindUsage {
+	lastChanged := u.UsedSlotsLastChanged
+	if other.UsedSlotsLastChanged.After(lastChanged) {
+		lastChanged = other.UsedSlotsLastChanged
+	}
+	return SlotKindUsage{
+		UsedSlots:            u.UsedSlots + other.UsedSlots,
+		UsedSlotsLastChanged: lastChanged,
+	}
+}
+
+// SlotUsage returns a snapshot of the worker's task slot usage. It is safe to call concurrently
+// with any other worker method, including before Start and after Stop. Task kinds whose workers
+// have not been created yet, such as Nexus before Start, report zero usage.
+func (aw *AggregatedWorker) SlotUsage() SlotUsage {
+	var u SlotUsage
+	if aw.workflowWorker != nil {
+		u.Workflow = aw.workflowWorker.worker.slotSupplier.usage()
+		u.LocalActivity = aw.workflowWorker.localActivityWorker.slotSupplier.usage()
+	}
+	if aw.activityWorker != nil {
+		u.Activity = u.Activity.add(aw.activityWorker.worker.slotSupplier.usage())
+	}
+	if aw.sessionWorker != nil {
+		u.Activity = u.Activity.add(aw.sessionWorker.activityWorker.worker.slotSupplier.usage())
+		u.Activity = u.Activity.add(aw.sessionWorker.creationWorker.worker.slotSupplier.usage())
+	}
+	if nexusSlots := aw.nexusSlotSupplier.Load(); nexusSlots != nil {
+		u.Nexus = nexusSlots.usage()
+	}
+	return u
+}

internal/slot_usage_test.go

@@ -0,0 +1,102 @@
+package internal
+
+import (
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.temporal.io/sdk/internal/common/metrics"
+	ilog "go.temporal.io/sdk/internal/log"
+)
+
+func newTestTrackingSlotSupplier(t *testing.T, numSlots int) *trackingSlotSupplier {
+	t.Helper()
+	inner, err := NewFixedSizeSlotSupplier(numSlots)
+	require.NoError(t, err)
+	return newTrackingSlotSupplier(inner, trackingSlotSupplierOptions{
+		logger:         ilog.NewNopLogger(),
+		metricsHandler: metrics.NopHandler,
+	})
+}
+
+func reserveTestSlot(t *testing.T, tss *trackingSlotSupplier) *SlotPermit {
+	t.Helper()
+	permit, err := tss.ReserveSlot(t.Context(), &slotReservationData{taskQueue: "test-task-queue"})
+	require.NoError(t, err)
+	return permit
+}
+
+func TestTrackingSlotSupplierUsage(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		tss := newTestTrackingSlotSupplier(t, 3)
+
+		// Reserved slots, like those held by idle pollers, are not in use.
+		used := reserveTestSlot(t, tss)
+		unused := reserveTestSlot(t, tss)
+		require.Equal(t, SlotKindUsage{}, tss.usage())
+
+		time.Sleep(time.Second)
+		startedAt := time.Now()
+		tss.MarkSlotUsed(used)
+		require.Equal(t, SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: startedAt}, tss.usage())
+
+		// Empty polls are released without being marked used, and must not count as a change.
+		time.Sleep(time.Second)
+		tss.ReleaseSlot(unused, SlotReleaseReasonTaskProcessed)
+		require.Equal(t, SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: startedAt}, tss.usage())
+
+		time.Sleep(time.Second)
+		finishedAt := time.Now()
+		tss.ReleaseSlot(used, SlotReleaseReasonTaskProcessed)
+		require.Equal(t, SlotKindUsage{UsedSlots: 0, UsedSlotsLastChanged: finishedAt}, tss.usage())
+	})
+}
+
+func TestAggregatedWorkerSlotUsage(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		require.Equal(t, SlotUsage{}, (&AggregatedWorker{}).SlotUsage())
+
+		workflowSlots := newTestTrackingSlotSupplier(t, 2)
+		localActivitySlots := newTestTrackingSlotSupplier(t, 2)
+		activitySlots := newTestTrackingSlotSupplier(t, 2)
+		sessionActivitySlots := newTestTrackingSlotSupplier(t, 2)
+		sessionCreationSlots := newTestTrackingSlotSupplier(t, 2)
+		nexusSlots := newTestTrackingSlotSupplier(t, 2)
+
+		aw := &AggregatedWorker{
+			workflowWorker: &workflowWorker{
+				worker:              &baseWorker{slotSupplier: workflowSlots},
+				localActivityWorker: &baseWorker{slotSupplier: localActivitySlots},
+			},
+			activityWorker: &activityWorker{worker: &baseWorker{slotSupplier: activitySlots}},
+			sessionWorker: &sessionWorker{
+				activityWorker: &activityWorker{worker: &baseWorker{slotSupplier: sessionActivitySlots}},
+				creationWorker: &activityWorker{worker: &baseWorker{slotSupplier: sessionCreationSlots}},
+			},
+		}
+
+		start := time.Now()
+		for _, tss := range []*trackingSlotSupplier{
+			workflowSlots, localActivitySlots, activitySlots, sessionActivitySlots, sessionCreationSlots, nexusSlots,
+		} {
+			tss.MarkSlotUsed(reserveTestSlot(t, tss))
+		}
+		// The latest change across the activity workers wins.
+		time.Sleep(time.Second)
+		finished := reserveTestSlot(t, activitySlots)
+		activitySlots.MarkSlotUsed(finished)
+		time.Sleep(time.Second)
+		activitySlots.ReleaseSlot(finished, SlotReleaseReasonTaskProcessed)
+
+		// The Nexus worker isn't created until start.
+		require.Equal(t, SlotUsage{
+			Workflow:      SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: start},
+			Activity:      SlotKindUsage{UsedSlots: 3, UsedSlotsLastChanged: start.Add(2 * time.Second)},
+			LocalActivity: SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: start},
+		}, aw.SlotUsage())
+
+		aw.nexusSlotSupplier.Store(nexusSlots)
+		require.Equal(t, SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: start}, aw.SlotUsage().Nexus)
+	})
+}

internal/tuning.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	enumspb "go.temporal.io/api/enums/v1"
 	"golang.org/x/sync/semaphore"
@@ -395,8 +396,10 @@ type trackingSlotSupplier struct {
 	issuedSlotsAtomic atomic.Int32
 	slotsMutex        sync.Mutex
 	// Values should eventually become slot info types
-	usedSlots               map[*SlotPermit]struct{}
-	usedSlotsVersion        uint64
+	usedSlots        map[*SlotPermit]struct{}
+	usedSlotsVersion uint64
+	// When a slot was last marked used or a used slot was released.
+	usedSlotsLastChanged    time.Time
 	taskSlotsAvailableGauge metrics.Gauge
 	taskSlotsUsedGauge      metrics.Gauge
 }
@@ -470,6 +473,7 @@ func (t *trackingSlotSupplier) MarkSlotUsed(permit *SlotPermit) {
 	t.slotsMutex.Lock()
 	t.usedSlots[permit] = struct{}{}
 	t.usedSlotsVersion++
+	t.usedSlotsLastChanged = time.Now()
 	t.slotsMutex.Unlock()
 	t.publishMetrics()
 	t.inner.MarkSlotUsed(&slotMarkUsedContextImpl{
@@ -484,7 +488,11 @@ func (t *trackingSlotSupplier) ReleaseSlot(permit *SlotPermit, reason SlotReleas
 		panic("Cannot release with nil permit")
 	}
 	t.slotsMutex.Lock()
-	delete(t.usedSlots, permit)
+	// Slots released without being used, such as after empty polls, don't change the usage.
+	if _, used := t.usedSlots[permit]; used {
+		delete(t.usedSlots, permit)
+		t.usedSlotsLastChanged = time.Now()
+	}
 	t.usedSlotsVersion++
 	t.slotsMutex.Unlock()
 	t.publishMetrics()
@@ -521,6 +529,15 @@ func (t *trackingSlotSupplier) publishMetrics() {
 	}
 }
 
+func (t *trackingSlotSupplier) usage() SlotKindUsage {
+	t.slotsMutex.Lock()
+	defer t.slotsMutex.Unlock()
+	return SlotKindUsage{
+		UsedSlots:            len(t.usedSlots),
+		UsedSlotsLastChanged: t.usedSlotsLastChanged,
+	}
+}
+
 func (t *trackingSlotSupplier) GetSlotSupplierKind() string {
 	return getSlotSupplierKind(t.inner)
 }

test/worker_tuner_test.go

@@ -7,8 +7,10 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	"go.temporal.io/sdk/activity"
 	"go.temporal.io/sdk/contrib/sysinfo"
 	"go.temporal.io/sdk/worker"
+	"go.temporal.io/sdk/workflow"
 )
 
 type WorkerTunerTestSuite struct {
@@ -138,6 +140,65 @@ func (ts *WorkerTunerTestSuite) TestResourceBasedSmallSlots() {
 	}
 }
 
+func (ts *WorkerTunerTestSuite) TestGetSlotUsage() {
+	ctx, cancel := context.WithTimeout(context.Background(), ctxTimeout)
+	defer cancel()
+
+	myWorker := worker.New(ts.client, ts.taskQueueName, worker.Options{})
+	activityStarted := make(chan struct{})
+	releaseActivity := make(chan struct{})
+	myWorker.RegisterActivityWithOptions(func(ctx context.Context) error {
+		close(activityStarted)
+		select {
+		case <-releaseActivity:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}, activity.RegisterOptions{Name: "SlotUsageBlockingActivity"})
+	myWorker.RegisterWorkflowWithOptions(func(ctx workflow.Context) error {
+		ctx = workflow.WithActivityOptions(ctx, workflow.ActivityOptions{StartToCloseTimeout: time.Minute})
+		return workflow.ExecuteActivity(ctx, "SlotUsageBlockingActivity").Get(ctx, nil)
+	}, workflow.RegisterOptions{Name: "SlotUsageWorkflow"})
+
+	usage, err := worker.GetSlotUsage(myWorker)
+	ts.NoError(err)
+	ts.Equal(worker.SlotUsage{}, usage)
+
+	ts.NoError(myWorker.Start())
+	defer myWorker.Stop()
+
+	startedAt := time.Now()
+	handle, err := ts.client.ExecuteWorkflow(ctx, ts.startWorkflowOptions(ts.T().Name()), "SlotUsageWorkflow")
+	ts.NoError(err)
+
+	select {
+	case <-activityStarted:
+	case <-ctx.Done():
+		ts.FailNow("activity did not start")
+	}
+	usage, err = worker.GetSlotUsage(myWorker)
+	ts.NoError(err)
+	ts.Equal(1, usage.Activity.UsedSlots)
+	ts.False(usage.Activity.UsedSlotsLastChanged.Before(startedAt))
+	ts.False(usage.Workflow.UsedSlotsLastChanged.Before(startedAt))
+
+	releasedAt := time.Now()
+	close(releaseActivity)
+	ts.NoError(handle.Get(ctx, nil))
+
+	// Slots are released asynchronously after task completions are sent to the server.
+	ts.Eventually(func() bool {
+		usage, err := worker.GetSlotUsage(myWorker)
+		return err == nil && usage.Workflow.UsedSlots == 0 && usage.Activity.UsedSlots == 0 &&
+			usage.LocalActivity.UsedSlots == 0 && usage.Nexus.UsedSlots == 0
+	}, 10*time.Second, 50*time.Millisecond)
+	usage, err = worker.GetSlotUsage(myWorker)
+	ts.NoError(err)
+	ts.False(usage.Activity.UsedSlotsLastChanged.Before(releasedAt))
+	ts.False(usage.Workflow.UsedSlotsLastChanged.Before(releasedAt))
+}
+
 func (ts *WorkerTunerTestSuite) runTheWorkflow(workerOptions worker.Options, ctx context.Context, wfTimeout time.Duration) {
 	myWorker := worker.New(ts.client, ts.taskQueueName, workerOptions)
 	ts.workflows.register(myWorker)

worker/slot_usage.go

@@ -0,0 +1,45 @@
+package worker
+
+import (
+	"fmt"
+
+	"go.temporal.io/sdk/internal"
+)
+
+// SlotUsage is a point-in-time snapshot of a worker's task slot usage, broken down by task kind.
+// Each task kind is read independently, so the snapshot is not atomic across task kinds.
+//
+// NOTE: Experimental
+type SlotUsage = internal.SlotUsage
+
+// SlotKindUsage is a snapshot of the task slot usage for one task kind.
+//
+// NOTE: Experimental
+type SlotKindUsage = internal.SlotKindUsage
+
+// SlotUsageProvider is implemented by workers that can report their task slot usage. Workers
+// created by [New] implement it. A type that wraps a worker can implement it by forwarding to the
+// wrapped worker.
+//
+// NOTE: Experimental
+type SlotUsageProvider interface {
+	// SlotUsage returns a snapshot of the worker's task slot usage. It must be safe to call
+	// concurrently with any other worker method.
+	SlotUsage() SlotUsage
+}
+
+var _ SlotUsageProvider = (*internal.AggregatedWorker)(nil)
+
+// GetSlotUsage returns a snapshot of w's task slot usage. It returns an error if w does not
+// implement [SlotUsageProvider].
+//
+// NOTE: Experimental
+func GetSlotUsage(w Worker) (SlotUsage, error) {
+	provider, ok := w.(SlotUsageProvider)
+	if !ok {
+		return SlotUsage{}, fmt.Errorf(
+			"worker of type %T does not implement worker.SlotUsageProvider; use a worker created by "+
+				"worker.New", w)
+	}
+	return provider.SlotUsage(), nil
+}