Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ProgressCheck callbacks to end-to-end acknowledgements #3565

Merged
merged 9 commits into from
Nov 4, 2023
Merged

Add ProgressCheck callbacks to end-to-end acknowledgements #3565

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
21f3526
Add ProgressCheck callbacks to end-to-end acknowledgements
Nov 2, 2023
5a9ab5c
Removed unnecessary logs
Nov 2, 2023
057b9d1
Fixed unnecessary change
Nov 2, 2023
211ac55
Fixed build failure in KafkaSourceJsonTypeIT
Nov 2, 2023
e8455cf
Fix sqs worker integration test failure
Nov 2, 2023
740cb4b
Addressed review comments
Nov 3, 2023
d3bfac4
Added Progress Check interface and test files
Nov 3, 2023
cf6e8f8
Addressed review comments
Nov 3, 2023
1b8fad6
Addressed review comments
Nov 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions ...i/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventHandle;

import java.time.Duration;
import java.util.function.Consumer;

/**
* AcknowledgmentSet keeps track of set of events that
* belong to the batch of events that a source creates.
Expand Down Expand Up @@ -58,4 +61,16 @@ public interface AcknowledgementSet {
* initial events are going through the pipeline line.
*/
public void complete();

/**
* adds progress check callback to the acknowledgement set. When added
* the callback is called every progressCheckInterval time with the
* indication of current progress as a ratio of pending number of
* acknowledgements over total acknowledgements
*
* @param progressCheckCallback progress check callback to be called
* @param progressCheckInterval frequency of invocation of progress check callback
* @since 2.6
*/
public void addProgressCheck(final Consumer<Double> progressCheckCallback, final Duration progressCheckInterval);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm curious why you chose to make this a Consumer<Double>. The current code isn't using this ratio presently.

A few alternatives:

  • Make this a Runnable without any input.
  • Create a new interface that can be provided so that the relevant arguments can be expanded in the future. For example, Consumer<ProgressCheck>.
public interface ProgressCheck {
  Double getRatio();
}

Then we could add other items if desired such as Duration getTimeToExpiry().

}
6 changes: 3 additions & 3 deletions ...e/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;

Expand All @@ -23,7 +23,7 @@ CallbackTheadFactory callbackTheadFactory() {
}

@Bean(name = "acknowledgementCallbackExecutor")
ExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) {
return Executors.newFixedThreadPool(MAX_THREADS, callbackTheadFactory);
ScheduledExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) {
return Executors.newScheduledThreadPool(MAX_THREADS, callbackTheadFactory);
}
}
46 changes: 40 additions & 6 deletions .../src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,37 +18,61 @@
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

public class DefaultAcknowledgementSet implements AcknowledgementSet {
private static final Logger LOG = LoggerFactory.getLogger(DefaultAcknowledgementSet.class);
private final Consumer<Boolean> callback;
private Consumer<Double> progressCheckCallback;
private final Instant expiryTime;
private final ExecutorService executor;
private final ScheduledExecutorService scheduledExecutor;
// This lock protects all the non-final members
private final ReentrantLock lock;
private boolean result;
private final Map<EventHandle, AtomicInteger> pendingAcknowledgments;
private Future<?> callbackFuture;
private final DefaultAcknowledgementSetMetrics metrics;
private ScheduledFuture<?> progressCheckFuture;
private boolean completed;
private AtomicInteger totalEventsAdded;

public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer<Boolean> callback, final Duration expiryTime, final DefaultAcknowledgementSetMetrics metrics) {
public DefaultAcknowledgementSet(final ScheduledExecutorService scheduledExecutor,
final Consumer<Boolean> callback,
final Duration expiryTime,
final DefaultAcknowledgementSetMetrics metrics) {
this.callback = callback;
this.result = true;
this.executor = executor;
this.totalEventsAdded = new AtomicInteger(0);
this.scheduledExecutor = scheduledExecutor;
this.expiryTime = Instant.now().plusMillis(expiryTime.toMillis());
this.callbackFuture = null;
this.metrics = metrics;
this.completed = false;
this.progressCheckCallback = null;
pendingAcknowledgments = new HashMap<>();
lock = new ReentrantLock(true);
}

public void addProgressCheck(final Consumer<Double> progressCheckCallback, final Duration progressCheckInterval) {
this.progressCheckCallback = progressCheckCallback;
this.progressCheckFuture = scheduledExecutor.scheduleAtFixedRate(this::checkProgress, 0L, progressCheckInterval.toMillis(), TimeUnit.MILLISECONDS);
}

public void checkProgress() {
lock.lock();
int numberOfEventsPending = pendingAcknowledgments.size();
lock.unlock();
if (progressCheckCallback != null) {
progressCheckCallback.accept((double)numberOfEventsPending/totalEventsAdded.get());
}
}

@Override
public void add(Event event) {
lock.lock();
Expand All @@ -59,6 +83,7 @@ public void add(Event event) {
InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
internalEventHandle.setAcknowledgementSet(this);
pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
totalEventsAdded.incrementAndGet();
}
}
} finally {
Expand Down Expand Up @@ -88,6 +113,9 @@ public boolean isDone() {
return true;
}
if (Instant.now().isAfter(expiryTime)) {
if (progressCheckFuture != null) {
progressCheckFuture.cancel(false);
}
if (callbackFuture != null) {
callbackFuture.cancel(true);
callbackFuture = null;
Expand All @@ -112,7 +140,10 @@ public void complete() {
try {
completed = true;
if (pendingAcknowledgments.size() == 0) {
callbackFuture = executor.submit(() -> callback.accept(this.result));
if (progressCheckFuture != null) {
progressCheckFuture.cancel(false);
}
callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result));
}
} finally {
lock.unlock();
Expand All @@ -136,7 +167,10 @@ public boolean release(final EventHandle eventHandle, final boolean result) {
if (pendingAcknowledgments.get(eventHandle).decrementAndGet() == 0) {
pendingAcknowledgments.remove(eventHandle);
if (completed && pendingAcknowledgments.size() == 0) {
callbackFuture = executor.submit(() -> callback.accept(this.result));
if (progressCheckFuture != null) {
progressCheckFuture.cancel(false);
}
callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result));
return true;
} else if (pendingAcknowledgments.size() == 0) {
LOG.warn("Acknowledgement set is not completed. Delaying callback until it is completed");
Expand Down
12 changes: 6 additions & 6 deletions ...in/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,35 +15,35 @@
import javax.inject.Named;
import java.time.Duration;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Consumer;

@Named
public class DefaultAcknowledgementSetManager implements AcknowledgementSetManager {
private static final int DEFAULT_WAIT_TIME_MS = 15 * 1000;
private final AcknowledgementSetMonitor acknowledgementSetMonitor;
private final ExecutorService executor;
private final ScheduledExecutorService scheduledExecutor;
private final AcknowledgementSetMonitorThread acknowledgementSetMonitorThread;
private PluginMetrics pluginMetrics;
private DefaultAcknowledgementSetMetrics metrics;

@Inject
public DefaultAcknowledgementSetManager(
@Named("acknowledgementCallbackExecutor") final ExecutorService callbackExecutor) {
@Named("acknowledgementCallbackExecutor") final ScheduledExecutorService callbackExecutor) {
this(callbackExecutor, Duration.ofMillis(DEFAULT_WAIT_TIME_MS));
}

public DefaultAcknowledgementSetManager(final ExecutorService callbackExecutor, final Duration waitTime) {
public DefaultAcknowledgementSetManager(final ScheduledExecutorService callbackExecutor, final Duration waitTime) {
this.acknowledgementSetMonitor = new AcknowledgementSetMonitor();
this.executor = Objects.requireNonNull(callbackExecutor);
this.scheduledExecutor = Objects.requireNonNull(callbackExecutor);
acknowledgementSetMonitorThread = new AcknowledgementSetMonitorThread(acknowledgementSetMonitor, waitTime);
acknowledgementSetMonitorThread.start();
pluginMetrics = PluginMetrics.fromNames("acknowledgementSetManager", "acknowledgements");
metrics = new DefaultAcknowledgementSetMetrics(pluginMetrics);
}

public AcknowledgementSet create(final Consumer<Boolean> callback, final Duration timeout) {
AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(executor, callback, timeout, metrics);
AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(scheduledExecutor, callback, timeout, metrics);
acknowledgementSetMonitor.add(acknowledgementSet);
metrics.increment(DefaultAcknowledgementSetMetrics.CREATED_METRIC_NAME);
return acknowledgementSet;
Expand Down
135 changes: 122 additions & 13 deletions ...va/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.metrics.PluginMetrics;

import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.junit.jupiter.api.BeforeEach;
Expand All @@ -22,14 +23,14 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Executors;

@ExtendWith(MockitoExtension.class)
class DefaultAcknowledgementSetManagerTests {
private static final Duration TEST_TIMEOUT = Duration.ofMillis(400);
DefaultAcknowledgementSetManager acknowledgementSetManager;
private ExecutorService callbackExecutor;
private DefaultAcknowledgementSetManager acknowledgementSetManager;
private ScheduledExecutorService callbackExecutor;

@Mock
JacksonEvent event1;
Expand All @@ -38,17 +39,24 @@ class DefaultAcknowledgementSetManagerTests {
@Mock
JacksonEvent event3;

DefaultEventHandle eventHandle1;
DefaultEventHandle eventHandle2;
DefaultEventHandle eventHandle3;
Boolean result;
private PluginMetrics pluginMetrics;
private DefaultEventHandle eventHandle1;
private DefaultEventHandle eventHandle2;
private DefaultEventHandle eventHandle3;
private DefaultEventHandle eventHandle4;
private DefaultEventHandle eventHandle5;
private DefaultEventHandle eventHandle6;
Comment on lines +43 to +48
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Could make this a List<DefaultEventHandle> just to make it more flexible for future test case additions

private Boolean result;
private double currentRatio;

@BeforeEach
void setup() {
callbackExecutor = Executors.newFixedThreadPool(2);
currentRatio = 0;
callbackExecutor = Executors.newScheduledThreadPool(2);
event1 = mock(JacksonEvent.class);
eventHandle1 = mock(DefaultEventHandle.class);
lenient().when(event1.getEventHandle()).thenReturn(eventHandle1);
pluginMetrics = mock(PluginMetrics.class);

event2 = mock(JacksonEvent.class);
eventHandle2 = mock(DefaultEventHandle.class);
Expand Down Expand Up @@ -76,16 +84,17 @@ void testBasic() {
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(true));
});
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(true));
}

@Test
void testExpirations() throws InterruptedException {
acknowledgementSetManager.releaseEventReference(eventHandle2, true);
Thread.sleep(TEST_TIMEOUT.multipliedBy(5).toMillis());
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(null));
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(result, equalTo(null));
});
}

@Test
Expand All @@ -106,7 +115,107 @@ void testMultipleAcknowledgementSets() {
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(true));
});
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(true));
}

@Test
void testWithProgressCheckCallbacks() {
eventHandle3 = mock(DefaultEventHandle.class);
lenient().when(event3.getEventHandle()).thenReturn(eventHandle3);

eventHandle4 = mock(DefaultEventHandle.class);
JacksonEvent event4 = mock(JacksonEvent.class);
lenient().when(event4.getEventHandle()).thenReturn(eventHandle4);

eventHandle5 = mock(DefaultEventHandle.class);
JacksonEvent event5 = mock(JacksonEvent.class);
lenient().when(event5.getEventHandle()).thenReturn(eventHandle5);

eventHandle6 = mock(DefaultEventHandle.class);
JacksonEvent event6 = mock(JacksonEvent.class);
lenient().when(event6.getEventHandle()).thenReturn(eventHandle6);

AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, Duration.ofMillis(10000));
acknowledgementSet2.addProgressCheck((ratio) -> {currentRatio = ratio;}, Duration.ofSeconds(1));
acknowledgementSet2.add(event3);
acknowledgementSet2.add(event4);
acknowledgementSet2.add(event5);
acknowledgementSet2.add(event6);
lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle4.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle5.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle6.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
acknowledgementSet2.complete();
acknowledgementSetManager.releaseEventReference(eventHandle3, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.75));
});
acknowledgementSetManager.releaseEventReference(eventHandle4, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.5));
});
acknowledgementSetManager.releaseEventReference(eventHandle5, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.25));
});
acknowledgementSetManager.releaseEventReference(eventHandle6, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(result, equalTo(true));
});

}

@Test
void testWithProgressCheckCallbacks_AcksExpire() {
eventHandle3 = mock(DefaultEventHandle.class);
lenient().when(event3.getEventHandle()).thenReturn(eventHandle3);

eventHandle4 = mock(DefaultEventHandle.class);
JacksonEvent event4 = mock(JacksonEvent.class);
lenient().when(event4.getEventHandle()).thenReturn(eventHandle4);

eventHandle5 = mock(DefaultEventHandle.class);
JacksonEvent event5 = mock(JacksonEvent.class);
lenient().when(event5.getEventHandle()).thenReturn(eventHandle5);

eventHandle6 = mock(DefaultEventHandle.class);
JacksonEvent event6 = mock(JacksonEvent.class);
lenient().when(event6.getEventHandle()).thenReturn(eventHandle6);

AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, Duration.ofSeconds(10));
acknowledgementSet2.addProgressCheck((ratio) -> {currentRatio = ratio;}, Duration.ofSeconds(1));
acknowledgementSet2.add(event3);
acknowledgementSet2.add(event4);
acknowledgementSet2.add(event5);
acknowledgementSet2.add(event6);
lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle4.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle5.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
lenient().when(eventHandle6.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
acknowledgementSet2.complete();
acknowledgementSetManager.releaseEventReference(eventHandle3, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.75));
});
acknowledgementSetManager.releaseEventReference(eventHandle4, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.5));
});
acknowledgementSetManager.releaseEventReference(eventHandle5, true);
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(currentRatio, equalTo(0.25));
});
await().atMost(TEST_TIMEOUT.multipliedBy(5))
.untilAsserted(() -> {
assertThat(result, equalTo(null));
});

}

}
Loading