Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 6dbf73d

Browse files
SushisourceSushisource
committed
Use a custom future type for slot suppliers
1 parent 8727900 commit 6dbf73d

File tree

12 files changed

+191
-120
lines changed

12 files changed

+191
-120
lines changed

‎temporal-sdk/src/main/java/io/temporal/internal/worker/ActivityPollTask.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
import io.temporal.worker.tuning.*;
3737
import java.util.Objects;
3838
import java.util.concurrent.ExecutionException;
39-
import java.util.concurrent.Future;
4039
import java.util.function.Supplier;
4140
import javax.annotation.Nonnull;
4241
import javax.annotation.Nullable;
@@ -98,7 +97,7 @@ public ActivityTask poll() {
9897
PollActivityTaskQueueResponse response;
9998
SlotPermit permit;
10099
boolean isSuccessful = false;
101-
Future<SlotPermit> future;
100+
SlotSupplierFuture future;
102101
try {
103102
future =
104103
slotSupplier.reserveSlot(
@@ -113,7 +112,7 @@ public ActivityTask poll() {
113112
try {
114113
permit = future.get();
115114
} catch (InterruptedException e) {
116-
future.cancel(true);
115+
future.abortReservation();
117116
Thread.currentThread().interrupt();
118117
return null;
119118
} catch (ExecutionException e) {

‎temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivitySlotSupplierQueue.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import io.temporal.worker.tuning.LocalActivitySlotInfo;
2424
import io.temporal.worker.tuning.SlotPermit;
2525
import io.temporal.worker.tuning.SlotReleaseReason;
26+
import io.temporal.worker.tuning.SlotSupplierFuture;
2627
import io.temporal.workflow.Functions;
2728
import java.util.concurrent.*;
2829
import javax.annotation.Nullable;
@@ -84,11 +85,11 @@ private void processQueue() {
8485
try {
8586
request = requestQueue.take();
8687

87-
Future<SlotPermit> future = slotSupplier.reserveSlot(request.data);
88+
SlotSupplierFuture future = slotSupplier.reserveSlot(request.data);
8889
try {
8990
slotPermit = future.get();
9091
} catch (InterruptedException e) {
91-
future.cancel(true);
92+
future.abortReservation();
9293
Thread.currentThread().interrupt();
9394
return;
9495
} catch (ExecutionException e) {

‎temporal-sdk/src/main/java/io/temporal/internal/worker/NexusPollTask.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import io.temporal.worker.tuning.*;
3434
import java.util.Objects;
3535
import java.util.concurrent.ExecutionException;
36-
import java.util.concurrent.Future;
3736
import java.util.function.Supplier;
3837
import javax.annotation.Nonnull;
3938
import javax.annotation.Nullable;
@@ -88,7 +87,7 @@ public NexusTask poll() {
8887
PollNexusTaskQueueResponse response;
8988
SlotPermit permit;
9089
boolean isSuccessful = false;
91-
Future<SlotPermit> future =
90+
SlotSupplierFuture future =
9291
slotSupplier.reserveSlot(
9392
new SlotReservationData(
9493
pollRequest.getTaskQueue().getName(),
@@ -97,7 +96,7 @@ public NexusTask poll() {
9796
try {
9897
permit = future.get();
9998
} catch (InterruptedException e) {
100-
future.cancel(true);
99+
future.abortReservation();
101100
Thread.currentThread().interrupt();
102101
return null;
103102
} catch (ExecutionException e) {

‎temporal-sdk/src/main/java/io/temporal/internal/worker/TrackingSlotSupplier.java‎

Lines changed: 7 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import java.util.Map;
2828
import java.util.Optional;
2929
import java.util.concurrent.*;
30-
import java.util.concurrent.atomic.AtomicBoolean;
3130
import java.util.concurrent.atomic.AtomicInteger;
3231

3332
/**
@@ -49,56 +48,20 @@ public TrackingSlotSupplier(SlotSupplier<SI> inner, Scope metricsScope) {
4948
publishSlotsMetric();
5049
}
5150

52-
public Future<SlotPermit> reserveSlot(SlotReservationData dat) {
53-
final Future<SlotPermit> originalFuture;
51+
public SlotSupplierFuture reserveSlot(SlotReservationData data) {
52+
final SlotSupplierFuture future;
5453
try {
55-
originalFuture = inner.reserveSlot(createCtx(dat));
54+
future = inner.reserveSlot(createCtx(data));
5655
} catch (Exception e) {
5756
throw new RuntimeException(e);
5857
}
5958

60-
return new Future<SlotPermit>() {
61-
private final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
62-
63-
private SlotPermit executeCallbackIfNeeded(SlotPermit permit) {
64-
if (callbackInvoked.compareAndSet(false, true)) {
65-
issuedSlots.incrementAndGet();
66-
}
67-
return permit;
68-
}
69-
70-
@Override
71-
public boolean cancel(boolean mayInterruptIfRunning) {
72-
return originalFuture.cancel(mayInterruptIfRunning);
73-
}
74-
75-
@Override
76-
public boolean isCancelled() {
77-
return originalFuture.isCancelled();
78-
}
79-
80-
@Override
81-
public boolean isDone() {
82-
return originalFuture.isDone();
83-
}
84-
85-
@Override
86-
public SlotPermit get() throws InterruptedException, ExecutionException {
87-
SlotPermit permit = originalFuture.get();
88-
return executeCallbackIfNeeded(permit);
89-
}
90-
91-
@Override
92-
public SlotPermit get(long timeout, TimeUnit unit)
93-
throws InterruptedException, ExecutionException, TimeoutException {
94-
SlotPermit permit = originalFuture.get(timeout, unit);
95-
return executeCallbackIfNeeded(permit);
96-
}
97-
};
59+
future.thenAccept(permit -> issuedSlots.incrementAndGet());
60+
return future;
9861
}
9962

100-
public Optional<SlotPermit> tryReserveSlot(SlotReservationData dat) {
101-
Optional<SlotPermit> p = inner.tryReserveSlot(createCtx(dat));
63+
public Optional<SlotPermit> tryReserveSlot(SlotReservationData data) {
64+
Optional<SlotPermit> p = inner.tryReserveSlot(createCtx(data));
10265
if (p.isPresent()) {
10366
issuedSlots.incrementAndGet();
10467
}

‎temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,10 @@
3737
import io.temporal.worker.MetricsType;
3838
import io.temporal.worker.tuning.SlotPermit;
3939
import io.temporal.worker.tuning.SlotReleaseReason;
40+
import io.temporal.worker.tuning.SlotSupplierFuture;
4041
import io.temporal.worker.tuning.WorkflowSlotInfo;
4142
import java.util.Objects;
4243
import java.util.concurrent.ExecutionException;
43-
import java.util.concurrent.Future;
4444
import java.util.function.Supplier;
4545
import javax.annotation.Nonnull;
4646
import javax.annotation.Nullable;
@@ -127,7 +127,7 @@ public WorkflowPollTask(
127127
public WorkflowTask poll() {
128128
boolean isSuccessful = false;
129129
SlotPermit permit;
130-
Future<SlotPermit> future =
130+
SlotSupplierFuture future =
131131
slotSupplier.reserveSlot(
132132
new SlotReservationData(
133133
pollRequest.getTaskQueue().getName(),
@@ -136,7 +136,7 @@ public WorkflowTask poll() {
136136
try {
137137
permit = future.get();
138138
} catch (InterruptedException e) {
139-
future.cancel(true);
139+
future.abortReservation();
140140
Thread.currentThread().interrupt();
141141
return null;
142142
} catch (ExecutionException e) {

‎temporal-sdk/src/main/java/io/temporal/worker/tuning/FixedSizeSlotSupplier.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import java.util.Optional;
2626
import java.util.Queue;
2727
import java.util.concurrent.CompletableFuture;
28-
import java.util.concurrent.Future;
2928
import java.util.concurrent.locks.ReentrantLock;
3029

3130
/**
@@ -94,7 +93,11 @@ public void release() {
9493
try {
9594
CompletableFuture<Void> waiter = waiters.poll();
9695
if (waiter != null) {
97-
waiter.complete(null);
96+
if (!waiter.complete(null) && waiter.isCancelled()) {
97+
// If this waiter was cancelled, we need to release another permit, since this waiter
98+
// is now useless
99+
release();
100+
}
98101
} else {
99102
permits++;
100103
}
@@ -111,8 +114,10 @@ public FixedSizeSlotSupplier(int numSlots) {
111114
}
112115

113116
@Override
114-
public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
115-
return executorSlotsSemaphore.acquire().thenApply(ignored -> new SlotPermit());
117+
public SlotSupplierFuture reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
118+
CompletableFuture<Void> slotFuture = executorSlotsSemaphore.acquire();
119+
return SlotSupplierFuture.fromCompletableFuture(
120+
slotFuture.thenApply(ignored -> new SlotPermit()), () -> slotFuture.cancel(true));
116121
}
117122

118123
@Override

‎temporal-sdk/src/main/java/io/temporal/worker/tuning/ResourceBasedSlotSupplier.java‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -202,16 +202,16 @@ private ResourceBasedSlotSupplier(
202202
}
203203

204204
@Override
205-
public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
205+
public SlotSupplierFuture reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
206206
if (ctx.getNumIssuedSlots() < options.getMinimumSlots()) {
207-
return CompletableFuture.completedFuture(new SlotPermit());
207+
return SlotSupplierFuture.completedFuture(new SlotPermit());
208208
}
209209
return tryReserveSlot(ctx)
210-
.map(CompletableFuture::completedFuture)
210+
.map(SlotSupplierFuture::completedFuture)
211211
.orElseGet(() -> scheduleSlotAcquisition(ctx));
212212
}
213213

214-
private CompletableFuture<SlotPermit> scheduleSlotAcquisition(SlotReserveContext<SI> ctx) {
214+
private SlotSupplierFuture scheduleSlotAcquisition(SlotReserveContext<SI> ctx) {
215215
Duration mustWaitFor;
216216
try {
217217
mustWaitFor = options.getRampThrottle().minus(timeSinceLastSlotIssued());
@@ -228,17 +228,19 @@ private CompletableFuture<SlotPermit> scheduleSlotAcquisition(SlotReserveContext
228228
}
229229

230230
// After the delay, try to reserve the slot
231-
return permitFuture.thenCompose(
232-
ignored -> {
233-
Optional<SlotPermit> permit = tryReserveSlot(ctx);
234-
// If we couldn't get a slot this time, delay for a short period and try again
235-
return permit
236-
.map(CompletableFuture::completedFuture)
237-
.orElseGet(
238-
() ->
239-
CompletableFuture.supplyAsync(() -> null, delayedExecutor(10))
240-
.thenCompose(ig -> scheduleSlotAcquisition(ctx)));
241-
});
231+
return SlotSupplierFuture.fromCompletableFuture(
232+
permitFuture.thenCompose(
233+
ignored -> {
234+
Optional<SlotPermit> permit = tryReserveSlot(ctx);
235+
// If we couldn't get a slot this time, delay for a short period and try again
236+
return permit
237+
.map(CompletableFuture::completedFuture)
238+
.orElseGet(
239+
() ->
240+
CompletableFuture.supplyAsync(() -> null, delayedExecutor(10))
241+
.thenCompose(ig -> scheduleSlotAcquisition(ctx)));
242+
}),
243+
() -> permitFuture.cancel(true));
242244
}
243245

244246
@Override

‎temporal-sdk/src/main/java/io/temporal/worker/tuning/SlotSupplier.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222

2323
import io.temporal.common.Experimental;
2424
import java.util.Optional;
25-
import java.util.concurrent.Future;
2625

2726
/**
2827
* A SlotSupplier is responsible for managing the number of slots available for a given type of
@@ -49,7 +48,7 @@ public interface SlotSupplier<SI extends SlotInfo> {
4948
* @return A future that will be completed with a permit to use the slot when one becomes
5049
* available. Never return null, or complete the future with null.
5150
*/
52-
Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception;
51+
SlotSupplierFuture reserveSlot(SlotReserveContext<SI> ctx) throws Exception;
5352

5453
/**
5554
* This function is called when trying to reserve slots for "eager" workflow and activity tasks.

‎temporal-sdk/src/main/java/io/temporal/worker/tuning/SlotSupplierFuture.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
/*
2+
* Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
3+
*
4+
* Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
*
6+
* Modifications copyright (C) 2017 Uber Technologies, Inc.
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this material except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
21+
package io.temporal.worker.tuning;
22+
23+
import java.util.concurrent.CompletableFuture;
24+
25+
/**
26+
* Represents a future that will be completed with a {@link SlotPermit} when a slot is available.
27+
*
28+
* <p>This class exists to provide a reliable cancellation mechanism, since {@link
29+
* CompletableFuture} does not provide cancellations that properly propagate up the chain.
30+
*/
31+
public abstract class SlotSupplierFuture extends CompletableFuture<SlotPermit> {
32+
/**
33+
* Abort the reservation attempt. Direct implementations should cancel or interrupt any underlying
34+
* processes that are attempting to reserve a slot.
35+
*/
36+
public abstract void abortReservation();
37+
38+
/** See {@link CompletableFuture#completedFuture(Object)} */
39+
public static SlotSupplierFuture completedFuture(SlotPermit permit) {
40+
return new SlotSupplierFuture() {
41+
@Override
42+
public void abortReservation() {}
43+
44+
{
45+
complete(permit);
46+
}
47+
};
48+
}
49+
50+
/**
51+
* Create a new {@link SlotSupplierFuture} from a {@link CompletableFuture}
52+
*
53+
* @param abortHandler The handler to call when the reservation is aborted. This should abort the
54+
* furthest-upstream future, or call being waited on, in order to properly propagate
55+
* cancellation downstream.
56+
*/
57+
public static SlotSupplierFuture fromCompletableFuture(
58+
CompletableFuture<SlotPermit> future, Runnable abortHandler) {
59+
SlotSupplierFuture wrapper =
60+
new SlotSupplierFuture() {
61+
@Override
62+
public void abortReservation() {
63+
abortHandler.run();
64+
future.cancel(true);
65+
}
66+
67+
@Override
68+
public boolean cancel(boolean mayInterruptIfRunning) {
69+
abortHandler.run();
70+
boolean cancelled = future.cancel(mayInterruptIfRunning);
71+
super.cancel(mayInterruptIfRunning);
72+
return cancelled;
73+
}
74+
};
75+
76+
future.whenComplete(
77+
(result, throwable) -> {
78+
if (throwable != null) {
79+
wrapper.completeExceptionally(throwable);
80+
} else {
81+
wrapper.complete(result);
82+
}
83+
});
84+
85+
return wrapper;
86+
}
87+
}

‎temporal-sdk/src/test/java/io/temporal/internal/worker/SlotSupplierTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@
3737
import io.temporal.serviceclient.WorkflowServiceStubs;
3838
import io.temporal.worker.tuning.*;
3939
import java.util.Objects;
40-
import java.util.concurrent.CompletableFuture;
4140
import java.util.concurrent.atomic.AtomicInteger;
4241
import org.junit.Test;
4342
import org.junit.runner.RunWith;
@@ -78,7 +77,7 @@ public void supplierIsCalledAppropriately() {
7877
usedSlotsWhenCalled.set(src.getUsedSlots().size());
7978
return true;
8079
})))
81-
.thenReturn(CompletableFuture.completedFuture(new SlotPermit()));
80+
.thenReturn(SlotSupplierFuture.completedFuture(new SlotPermit()));
8281
} catch (Exception e) {
8382
throw new RuntimeException(e);
8483
}

0 commit comments

Comments
 (0)