diff --git a/CHANGELOG.md b/CHANGELOG.md index 03a7ea695..f57d2a9ad 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/internal/internal_worker.go b/internal/internal_worker.go index 0469fc152..a8f927d36 100644 --- a/internal/internal_worker.go +++ b/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) } diff --git a/internal/slot_usage.go b/internal/slot_usage.go new file mode 100644 index 000000000..6ef5ad8ad --- /dev/null +++ b/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 +} diff --git a/internal/slot_usage_test.go b/internal/slot_usage_test.go new file mode 100644 index 000000000..dea9ed48a --- /dev/null +++ b/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) + }) +} diff --git a/internal/tuning.go b/internal/tuning.go index 25b740714..e95bcacae 100644 --- a/internal/tuning.go +++ b/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) } diff --git a/test/worker_tuner_test.go b/test/worker_tuner_test.go index aaa5325c4..d3103dc10 100644 --- a/test/worker_tuner_test.go +++ b/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) diff --git a/worker/slot_usage.go b/worker/slot_usage.go new file mode 100644 index 000000000..3ab19033f --- /dev/null +++ b/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 +} diff --git a/worker/slot_usage_test.go b/worker/slot_usage_test.go new file mode 100644 index 000000000..a45fe4287 --- /dev/null +++ b/worker/slot_usage_test.go @@ -0,0 +1,35 @@ +package worker_test + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.temporal.io/sdk/worker" +) + +type workerWithoutSlotUsage struct { + worker.Worker +} + +type workerWithSlotUsage struct { + worker.Worker + usage worker.SlotUsage +} + +func (w workerWithSlotUsage) SlotUsage() worker.SlotUsage { + return w.usage +} + +func TestGetSlotUsage(t *testing.T) { + _, err := worker.GetSlotUsage(workerWithoutSlotUsage{}) + require.ErrorContains(t, err, "does not implement worker.SlotUsageProvider") + + _, err = worker.GetSlotUsage(nil) + require.ErrorContains(t, err, "does not implement worker.SlotUsageProvider") + + want := worker.SlotUsage{Activity: worker.SlotKindUsage{UsedSlots: 1, UsedSlotsLastChanged: time.Unix(1, 0)}} + got, err := worker.GetSlotUsage(workerWithSlotUsage{usage: want}) + require.NoError(t, err) + require.Equal(t, want, got) +}