diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java index a34e55d904..f6692e2144 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java @@ -321,6 +321,9 @@ public String getBinaryChecksum() { } public String getBuildId() { + if (deploymentOptions != null && deploymentOptions.getVersion() != null) { + return deploymentOptions.getVersion().getBuildId(); + } if (buildId == null) { return binaryChecksum; } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java index 1b6c8cf7dc..f9c3e2f103 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java @@ -163,7 +163,7 @@ public WorkflowTask poll() { isSuccessful = true; tracker.pollSucceeded(); stickyQueueBalancer.finishPoll(taskQueueKind, response.getBacklogCountHint()); - slotSupplier.markSlotUsed(new WorkflowSlotInfo(response, pollRequest), permit); + slotSupplier.markSlotUsed(new WorkflowSlotInfo(response, request), permit); return new WorkflowTask(response, (rr) -> slotSupplier.releaseSlot(rr, permit)); } finally { if (isSticky) { diff --git a/temporal-sdk/src/main/java/io/temporal/worker/tuning/WorkflowSlotInfo.java b/temporal-sdk/src/main/java/io/temporal/worker/tuning/WorkflowSlotInfo.java index 9310bc64cf..e0afe31d11 100644 --- a/temporal-sdk/src/main/java/io/temporal/worker/tuning/WorkflowSlotInfo.java +++ b/temporal-sdk/src/main/java/io/temporal/worker/tuning/WorkflowSlotInfo.java @@ -22,11 +22,20 @@ public WorkflowSlotInfo( @Nonnull PollWorkflowTaskQueueResponse response, @Nonnull PollWorkflowTaskQueueRequest request) { this.workflowType = response.getWorkflowType().getName(); - this.taskQueue = request.getTaskQueue().getNormalName(); + this.taskQueue = + request.getTaskQueue().getKind() == TaskQueueKind.TASK_QUEUE_KIND_STICKY + ? request.getTaskQueue().getNormalName() + : request.getTaskQueue().getName(); this.workflowId = response.getWorkflowExecution().getWorkflowId(); this.runId = response.getWorkflowExecution().getRunId(); this.workerIdentity = request.getIdentity(); - this.workerBuildId = request.getWorkerVersionCapabilities().getBuildId(); + if (request.hasDeploymentOptions()) { + this.workerBuildId = request.getDeploymentOptions().getBuildId(); + } else if (request.hasWorkerVersionCapabilities()) { + this.workerBuildId = request.getWorkerVersionCapabilities().getBuildId(); + } else { + this.workerBuildId = request.getBinaryChecksum(); + } this.fromStickyQueue = request.getTaskQueue().getKind() == TaskQueueKind.TASK_QUEUE_KIND_STICKY; } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/SlotInfoTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/SlotInfoTest.java new file mode 100644 index 0000000000..02054b8a5c --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/SlotInfoTest.java @@ -0,0 +1,268 @@ +package io.temporal.internal.worker; + +import static io.temporal.testUtils.Eventually.assertEventually; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertSame; + +import io.nexusrpc.handler.OperationHandler; +import io.nexusrpc.handler.OperationImpl; +import io.nexusrpc.handler.ServiceImpl; +import io.temporal.activity.ActivityInterface; +import io.temporal.activity.ActivityMethod; +import io.temporal.activity.ActivityOptions; +import io.temporal.activity.LocalActivityOptions; +import io.temporal.client.WorkflowOptions; +import io.temporal.client.WorkflowStub; +import io.temporal.common.RetryOptions; +import io.temporal.common.WorkerDeploymentVersion; +import io.temporal.testUtils.RecordingSlotSupplier; +import io.temporal.testing.internal.SDKTestWorkflowRule; +import io.temporal.worker.WorkerDeploymentOptions; +import io.temporal.worker.WorkerOptions; +import io.temporal.worker.tuning.ActivitySlotInfo; +import io.temporal.worker.tuning.CompositeTuner; +import io.temporal.worker.tuning.LocalActivitySlotInfo; +import io.temporal.worker.tuning.NexusSlotInfo; +import io.temporal.worker.tuning.SlotInfo; +import io.temporal.worker.tuning.SlotMarkUsedContext; +import io.temporal.worker.tuning.SlotPermit; +import io.temporal.worker.tuning.SlotReleaseContext; +import io.temporal.worker.tuning.WorkflowSlotInfo; +import io.temporal.workflow.NexusOperationOptions; +import io.temporal.workflow.NexusServiceOptions; +import io.temporal.workflow.Workflow; +import io.temporal.workflow.WorkflowInterface; +import io.temporal.workflow.WorkflowMethod; +import io.temporal.workflow.shared.TestNexusServices; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import org.junit.Rule; +import org.junit.Test; + +@SuppressWarnings("deprecation") +public class SlotInfoTest { + private static final String WORKFLOW_TYPE = "slot-info-workflow"; + private static final String ACTIVITY_TYPE = "slot-info-activity"; + private static final String WORKFLOW_ID = "slot-info-workflow-id"; + private static final String WORKER_IDENTITY = "slot-info-worker-identity"; + private static final String WORKER_BUILD_ID = "slot-info-worker-build-id"; + + private final RecordingSlotSupplier workflowSlotSupplier = + new RecordingSlotSupplier<>(100); + private final RecordingSlotSupplier activitySlotSupplier = + new RecordingSlotSupplier<>(100); + private final RecordingSlotSupplier localActivitySlotSupplier = + new RecordingSlotSupplier<>(100); + private final RecordingSlotSupplier nexusSlotSupplier = + new RecordingSlotSupplier<>(100); + + @Rule + public SDKTestWorkflowRule testWorkflowRule = + SDKTestWorkflowRule.newBuilder() + .setWorkerOptions( + WorkerOptions.newBuilder() + .setIdentity(WORKER_IDENTITY) + .setDeploymentOptions( + WorkerDeploymentOptions.newBuilder() + .setVersion( + new WorkerDeploymentVersion("slot-info-deployment", WORKER_BUILD_ID)) + .setUseVersioning(false) + .build()) + .setWorkerTuner( + new CompositeTuner( + workflowSlotSupplier, + activitySlotSupplier, + localActivitySlotSupplier, + nexusSlotSupplier)) + .build()) + .setWorkflowTypes(SlotInfoWorkflowImpl.class) + .setActivityImplementations(new SlotInfoActivityImpl()) + .setNexusServiceImplementation(new SlotInfoNexusService()) + .build(); + + @Test + public void customSlotSuppliersReceiveExpectedSlotInfo() { + SlotInfoWorkflow workflow = + testWorkflowRule + .getWorkflowClient() + .newWorkflowStub( + SlotInfoWorkflow.class, + WorkflowOptions.newBuilder() + .setWorkflowId(WORKFLOW_ID) + .setTaskQueue(testWorkflowRule.getTaskQueue()) + .build()); + + assertEquals("done", workflow.execute()); + String runId = WorkflowStub.fromTyped(workflow).getExecution().getRunId(); + + List workflowSlotInfos = getSlotInfos(workflowSlotSupplier); + assertFalse(workflowSlotInfos.isEmpty()); + for (WorkflowSlotInfo slotInfo : workflowSlotInfos) { + assertEquals(WORKFLOW_TYPE, slotInfo.getWorkflowType()); + assertEquals(testWorkflowRule.getTaskQueue(), slotInfo.getTaskQueue()); + assertEquals(WORKFLOW_ID, slotInfo.getWorkflowId()); + assertEquals(runId, slotInfo.getRunId()); + assertEquals(WORKER_IDENTITY, slotInfo.getWorkerIdentity()); + assertEquals(WORKER_BUILD_ID, slotInfo.getWorkerBuildId()); + } + + ActivitySlotInfo activitySlotInfo = getOnlySlotInfo(activitySlotSupplier); + assertActivityInfo(activitySlotInfo, runId, false); + + LocalActivitySlotInfo localActivitySlotInfo = getOnlySlotInfo(localActivitySlotSupplier); + assertActivityInfo(localActivitySlotInfo, runId, true); + + NexusSlotInfo nexusSlotInfo = getOnlySlotInfo(nexusSlotSupplier); + assertEquals( + TestNexusServices.TestNexusService1.class.getSimpleName(), nexusSlotInfo.getService()); + assertEquals("operation", nexusSlotInfo.getOperation()); + assertEquals(testWorkflowRule.getTaskQueue(), nexusSlotInfo.getTaskQueue()); + assertEquals(WORKER_IDENTITY, nexusSlotInfo.getWorkerIdentity()); + assertEquals(WORKER_BUILD_ID, nexusSlotInfo.getWorkerBuildId()); + + workflowSlotInfos.forEach( + slotInfo -> assertReleasedWithSameSlotInfo(workflowSlotSupplier, slotInfo)); + assertReleasedWithSameSlotInfo(activitySlotSupplier, activitySlotInfo); + assertReleasedWithSameSlotInfo(localActivitySlotSupplier, localActivitySlotInfo); + assertReleasedWithSameSlotInfo(nexusSlotSupplier, nexusSlotInfo); + } + + private static List getSlotInfos( + RecordingSlotSupplier slotSupplier) { + List result = new ArrayList<>(); + for (SlotMarkUsedContext context : slotSupplier.getMarkUsedContexts()) { + result.add(context.getSlotInfo()); + } + return result; + } + + private static SI getOnlySlotInfo(RecordingSlotSupplier slotSupplier) { + List slotInfos = getSlotInfos(slotSupplier); + assertEquals(1, slotInfos.size()); + return slotInfos.get(0); + } + + private static void assertReleasedWithSameSlotInfo( + RecordingSlotSupplier slotSupplier, SI expectedSlotInfo) { + SlotPermit permit = null; + for (SlotMarkUsedContext context : slotSupplier.getMarkUsedContexts()) { + if (context.getSlotInfo() == expectedSlotInfo) { + permit = context.getSlotPermit(); + break; + } + } + assertNotNull(permit); + SlotPermit markedPermit = permit; + assertEventually( + Duration.ofSeconds(1), + () -> { + SI releasedSlotInfo = null; + for (SlotReleaseContext context : slotSupplier.getReleaseContexts()) { + if (context.getSlotPermit() == markedPermit) { + releasedSlotInfo = context.getSlotInfo(); + break; + } + } + assertSame(expectedSlotInfo, releasedSlotInfo); + }); + } + + private void assertActivityInfo(ActivitySlotInfo slotInfo, String runId, boolean expectedLocal) { + assertActivityInfo( + slotInfo.getActivityInfo(), + slotInfo.getWorkerIdentity(), + slotInfo.getWorkerBuildId(), + runId, + expectedLocal); + } + + private void assertActivityInfo( + LocalActivitySlotInfo slotInfo, String runId, boolean expectedLocal) { + assertActivityInfo( + slotInfo.getActivityInfo(), + slotInfo.getWorkerIdentity(), + slotInfo.getWorkerBuildId(), + runId, + expectedLocal); + } + + private void assertActivityInfo( + io.temporal.activity.ActivityInfo activityInfo, + String workerIdentity, + String workerBuildId, + String runId, + boolean expectedLocal) { + assertEquals(ACTIVITY_TYPE, activityInfo.getActivityType()); + assertFalse(activityInfo.getActivityId().isEmpty()); + assertEquals(WORKFLOW_ID, activityInfo.getWorkflowId()); + assertEquals(runId, activityInfo.getWorkflowRunId()); + assertEquals(WORKFLOW_TYPE, activityInfo.getWorkflowType()); + assertEquals(testWorkflowRule.getTaskQueue(), activityInfo.getActivityTaskQueue()); + assertEquals(SDKTestWorkflowRule.NAMESPACE, activityInfo.getNamespace()); + assertEquals(1, activityInfo.getAttempt()); + assertEquals(expectedLocal, activityInfo.isLocal()); + assertEquals(WORKER_IDENTITY, workerIdentity); + assertEquals(WORKER_BUILD_ID, workerBuildId); + } + + @WorkflowInterface + public interface SlotInfoWorkflow { + @WorkflowMethod(name = WORKFLOW_TYPE) + String execute(); + } + + public static class SlotInfoWorkflowImpl implements SlotInfoWorkflow { + private final SlotInfoActivity activity = + Workflow.newActivityStub( + SlotInfoActivity.class, + ActivityOptions.newBuilder().setStartToCloseTimeout(Duration.ofSeconds(10)).build()); + private final SlotInfoActivity localActivity = + Workflow.newLocalActivityStub( + SlotInfoActivity.class, + LocalActivityOptions.newBuilder() + .setStartToCloseTimeout(Duration.ofSeconds(10)) + .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(1).build()) + .build()); + private final TestNexusServices.TestNexusService1 nexusService = + Workflow.newNexusServiceStub( + TestNexusServices.TestNexusService1.class, + NexusServiceOptions.newBuilder() + .setOperationOptions( + NexusOperationOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofSeconds(10)) + .build()) + .build()); + + @Override + public String execute() { + localActivity.execute(); + activity.execute(); + nexusService.operation("input"); + return "done"; + } + } + + @ActivityInterface + public interface SlotInfoActivity { + @ActivityMethod(name = ACTIVITY_TYPE) + String execute(); + } + + public static class SlotInfoActivityImpl implements SlotInfoActivity { + @Override + public String execute() { + return "done"; + } + } + + @ServiceImpl(service = TestNexusServices.TestNexusService1.class) + public static class SlotInfoNexusService { + @OperationImpl + public OperationHandler operation() { + return OperationHandler.sync((ctx, details, input) -> "done"); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowSlotInfoTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowSlotInfoTest.java new file mode 100644 index 0000000000..b4679f3fcb --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowSlotInfoTest.java @@ -0,0 +1,227 @@ +package io.temporal.internal.worker; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.util.concurrent.Futures; +import com.google.protobuf.ByteString; +import com.uber.m3.tally.NoopScope; +import io.temporal.api.common.v1.WorkerVersionCapabilities; +import io.temporal.api.common.v1.WorkflowExecution; +import io.temporal.api.common.v1.WorkflowType; +import io.temporal.api.enums.v1.TaskQueueKind; +import io.temporal.api.taskqueue.v1.TaskQueue; +import io.temporal.api.workflowservice.v1.GetSystemInfoResponse; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueRequest; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.testUtils.RecordingSlotSupplier; +import io.temporal.worker.tuning.SlotInfo; +import io.temporal.worker.tuning.SlotMarkUsedContext; +import io.temporal.worker.tuning.SlotPermit; +import io.temporal.worker.tuning.SlotReleaseContext; +import io.temporal.worker.tuning.WorkflowSlotInfo; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import org.junit.Test; + +@SuppressWarnings("deprecation") +public class WorkflowSlotInfoTest { + private static final String WORKFLOW_TYPE = "test-workflow-type"; + private static final String TASK_QUEUE = "test-task-queue"; + private static final String STICKY_TASK_QUEUE = "test-sticky-task-queue"; + private static final String WORKFLOW_ID = "test-workflow-id"; + private static final String RUN_ID = "test-run-id"; + private static final String WORKER_IDENTITY = "test-worker-identity"; + private static final String WORKER_BUILD_ID = "test-worker-build-id"; + + @Test + public void normalWorkflowSlotInfoHasExpectedFields() { + PollWorkflowTaskQueueRequest request = + PollWorkflowTaskQueueRequest.newBuilder() + .setIdentity(WORKER_IDENTITY) + .setTaskQueue( + TaskQueue.newBuilder() + .setName(TASK_QUEUE) + .setKind(TaskQueueKind.TASK_QUEUE_KIND_NORMAL)) + .setWorkerVersionCapabilities( + WorkerVersionCapabilities.newBuilder().setBuildId(WORKER_BUILD_ID)) + .build(); + + WorkflowSlotInfo slotInfo = new WorkflowSlotInfo(workflowResponse(), request); + + assertWorkflowSlotInfo(slotInfo, false); + } + + @Test + public void deploymentBuildIdIsIncludedInWorkflowSlotInfo() { + PollWorkflowTaskQueueRequest request = + normalPollRequestBuilder() + .setDeploymentOptions( + io.temporal.api.deployment.v1.WorkerDeploymentOptions.newBuilder() + .setBuildId(WORKER_BUILD_ID)) + .build(); + + WorkflowSlotInfo slotInfo = new WorkflowSlotInfo(workflowResponse(), request); + + assertEquals(WORKER_BUILD_ID, slotInfo.getWorkerBuildId()); + } + + @Test + public void binaryChecksumIsIncludedInWorkflowSlotInfo() { + PollWorkflowTaskQueueRequest request = + normalPollRequestBuilder().setBinaryChecksum(WORKER_BUILD_ID).build(); + + WorkflowSlotInfo slotInfo = new WorkflowSlotInfo(workflowResponse(), request); + + assertEquals(WORKER_BUILD_ID, slotInfo.getWorkerBuildId()); + } + + @Test + public void synchronousStickyPollUsesSelectedRequestForSlotInfo() { + WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); + WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub = + mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class); + when(client.blockingStub()).thenReturn(blockingStub); + when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); + when(blockingStub.pollWorkflowTaskQueue(any())).thenReturn(workflowResponse()); + + RecordingSlotSupplier recordingSupplier = new RecordingSlotSupplier<>(1); + TrackingSlotSupplier trackingSupplier = + new TrackingSlotSupplier<>(recordingSupplier, new NoopScope()); + WorkflowPollTask pollTask = + new WorkflowPollTask( + client, + "default", + TASK_QUEUE, + STICKY_TASK_QUEUE, + WORKER_IDENTITY, + "test-instance-key", + new WorkerVersioningOptions(WORKER_BUILD_ID, false, null), + trackingSupplier, + new StickyQueueBalancer(1, true), + new NoopScope(), + WorkflowSlotInfoTest::buildIdCapabilities, + new PollerTracker(), + new PollerTracker(), + null); + + WorkflowTask task = pollTask.poll(); + + assertNotNull(task); + SlotMarkUsedContext markUsedContext = + getOnlyMarkUsedContext(recordingSupplier); + List reservedPermits = recordingSupplier.getReservedPermits(); + assertEquals(1, reservedPermits.size()); + assertSame(reservedPermits.get(0), markUsedContext.getSlotPermit()); + assertWorkflowSlotInfo(markUsedContext.getSlotInfo(), true); + + task.getCompletionCallback().apply(io.temporal.worker.tuning.SlotReleaseReason.taskComplete()); + SlotReleaseContext releaseContext = getOnlyReleaseContext(recordingSupplier); + assertSame(markUsedContext.getSlotPermit(), releaseContext.getSlotPermit()); + assertSame(markUsedContext.getSlotInfo(), releaseContext.getSlotInfo()); + } + + @Test + public void asynchronousPollsIncludeNormalAndStickyQueueFields() throws Exception { + assertAsyncWorkflowSlotInfo(null, false); + assertAsyncWorkflowSlotInfo(STICKY_TASK_QUEUE, true); + } + + private static void assertAsyncWorkflowSlotInfo(String stickyTaskQueue, boolean expectedSticky) + throws Exception { + WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); + WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = + mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); + when(client.futureStub()).thenReturn(futureStub); + when(futureStub.withOption(any(), any())).thenReturn(futureStub); + when(futureStub.pollWorkflowTaskQueue(any())) + .thenReturn(Futures.immediateFuture(workflowResponse())); + + RecordingSlotSupplier recordingSupplier = new RecordingSlotSupplier<>(1); + TrackingSlotSupplier trackingSupplier = + new TrackingSlotSupplier<>(recordingSupplier, new NoopScope()); + AsyncWorkflowPollTask pollTask = + new AsyncWorkflowPollTask( + client, + "default", + TASK_QUEUE, + stickyTaskQueue, + WORKER_IDENTITY, + "test-instance-key", + new WorkerVersioningOptions(WORKER_BUILD_ID, false, null), + trackingSupplier, + new NoopScope(), + WorkflowSlotInfoTest::buildIdCapabilities, + new PollerTracker(), + null); + SlotPermit permit = new SlotPermit(); + + CompletableFuture future = pollTask.poll(permit); + WorkflowTask task = future.get(); + + assertNotNull(task); + SlotMarkUsedContext markUsedContext = + getOnlyMarkUsedContext(recordingSupplier); + assertSame(permit, markUsedContext.getSlotPermit()); + assertWorkflowSlotInfo(markUsedContext.getSlotInfo(), expectedSticky); + } + + private static SlotMarkUsedContext getOnlyMarkUsedContext( + RecordingSlotSupplier slotSupplier) { + List> contexts = slotSupplier.getMarkUsedContexts(); + assertEquals(1, contexts.size()); + return contexts.get(0); + } + + private static SlotReleaseContext getOnlyReleaseContext( + RecordingSlotSupplier slotSupplier) { + List> contexts = slotSupplier.getReleaseContexts(); + assertEquals(1, contexts.size()); + return contexts.get(0); + } + + private static PollWorkflowTaskQueueRequest.Builder normalPollRequestBuilder() { + return PollWorkflowTaskQueueRequest.newBuilder() + .setIdentity(WORKER_IDENTITY) + .setTaskQueue( + TaskQueue.newBuilder() + .setName(TASK_QUEUE) + .setKind(TaskQueueKind.TASK_QUEUE_KIND_NORMAL)); + } + + private static PollWorkflowTaskQueueResponse workflowResponse() { + return PollWorkflowTaskQueueResponse.newBuilder() + .setTaskToken(ByteString.copyFrom("token", UTF_8)) + .setWorkflowExecution( + WorkflowExecution.newBuilder().setWorkflowId(WORKFLOW_ID).setRunId(RUN_ID)) + .setWorkflowType(WorkflowType.newBuilder().setName(WORKFLOW_TYPE)) + .build(); + } + + private static GetSystemInfoResponse.Capabilities buildIdCapabilities() { + return GetSystemInfoResponse.Capabilities.newBuilder().setBuildIdBasedVersioning(true).build(); + } + + private static void assertWorkflowSlotInfo(WorkflowSlotInfo slotInfo, boolean expectedSticky) { + assertEquals(WORKFLOW_TYPE, slotInfo.getWorkflowType()); + assertEquals(TASK_QUEUE, slotInfo.getTaskQueue()); + assertEquals(WORKFLOW_ID, slotInfo.getWorkflowId()); + assertEquals(RUN_ID, slotInfo.getRunId()); + assertEquals(WORKER_IDENTITY, slotInfo.getWorkerIdentity()); + assertEquals(WORKER_BUILD_ID, slotInfo.getWorkerBuildId()); + if (expectedSticky) { + assertTrue(slotInfo.isFromStickyQueue()); + } else { + assertFalse(slotInfo.isFromStickyQueue()); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/testUtils/RecordingSlotSupplier.java b/temporal-sdk/src/test/java/io/temporal/testUtils/RecordingSlotSupplier.java new file mode 100644 index 0000000000..adeec204d9 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/testUtils/RecordingSlotSupplier.java @@ -0,0 +1,64 @@ +package io.temporal.testUtils; + +import io.temporal.worker.tuning.FixedSizeSlotSupplier; +import io.temporal.worker.tuning.SlotInfo; +import io.temporal.worker.tuning.SlotMarkUsedContext; +import io.temporal.worker.tuning.SlotPermit; +import io.temporal.worker.tuning.SlotReleaseContext; +import io.temporal.worker.tuning.SlotReserveContext; +import io.temporal.worker.tuning.SlotSupplierFuture; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.ConcurrentLinkedQueue; + +/** A fixed-size slot supplier that records slot usage for tests. */ +public final class RecordingSlotSupplier extends FixedSizeSlotSupplier { + private final ConcurrentLinkedQueue reservedPermits = new ConcurrentLinkedQueue<>(); + private final ConcurrentLinkedQueue> markUsedContexts = + new ConcurrentLinkedQueue<>(); + private final ConcurrentLinkedQueue> releaseContexts = + new ConcurrentLinkedQueue<>(); + + public RecordingSlotSupplier(int numSlots) { + super(numSlots); + } + + @Override + public SlotSupplierFuture reserveSlot(SlotReserveContext ctx) throws Exception { + SlotSupplierFuture future = super.reserveSlot(ctx); + future.thenAccept(reservedPermits::add); + return future; + } + + @Override + public Optional tryReserveSlot(SlotReserveContext ctx) { + Optional permit = super.tryReserveSlot(ctx); + permit.ifPresent(reservedPermits::add); + return permit; + } + + @Override + public void markSlotUsed(SlotMarkUsedContext ctx) { + markUsedContexts.add(ctx); + super.markSlotUsed(ctx); + } + + @Override + public void releaseSlot(SlotReleaseContext ctx) { + releaseContexts.add(ctx); + super.releaseSlot(ctx); + } + + public List> getMarkUsedContexts() { + return new ArrayList<>(markUsedContexts); + } + + public List getReservedPermits() { + return new ArrayList<>(reservedPermits); + } + + public List> getReleaseContexts() { + return new ArrayList<>(releaseContexts); + } +}