> {
* @throws RuntimeException Other exceptions
*/
default void writeBytes(final byte[] bytes, final String key, int timeoutInMillis) throws Exception {
- throw new RuntimeException("Not supported");
+ throw new UnsupportedOperationException("This buffer type does not support bytes.");
}
/**
@@ -92,6 +92,17 @@ default Duration getDrainTimeout() {
return Duration.ZERO;
}
+ /**
+ * Indicates if writes to this buffer are also in some way written
+ * onto the JVM heap. If writes do go on heap, this should false
+ * which is the default.
+ *
+ * @return True if this buffer does not write to the JVM heap.
+ */
+ default boolean isWrittenOffHeapOnly() {
+ return false;
+ }
+
/**
* shuts down the buffer
*/
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java
new file mode 100644
index 0000000000..84f57883ea
--- /dev/null
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.buffer;
+
+import org.opensearch.dataprepper.model.CheckpointState;
+import org.opensearch.dataprepper.model.record.Record;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * An implementation of {@link Buffer} which delegates all calls to a delgate
+ * (or inner) buffer.
+ *
+ * This class exists to help with writing decorators of the {@link Buffer} interface.
+ *
+ * @param The type of data in the buffer
+ *
+ * @since 2.6
+ */
+public abstract class DelegatingBuffer> implements Buffer {
+ private final Buffer delegateBuffer;
+
+ /**
+ * Constructor for subclasses to use.
+ *
+ * @param delegateBuffer The delegate (or inner) buffer.
+ *
+ * @since 2.6
+ */
+ protected DelegatingBuffer(final Buffer delegateBuffer) {
+ this.delegateBuffer = Objects.requireNonNull(delegateBuffer);
+ }
+
+ @Override
+ public void write(final T record, final int timeoutInMillis) throws TimeoutException {
+ delegateBuffer.write(record, timeoutInMillis);
+ }
+
+ @Override
+ public void writeAll(final Collection records, final int timeoutInMillis) throws Exception {
+ delegateBuffer.writeAll(records, timeoutInMillis);
+ }
+
+ @Override
+ public void writeBytes(final byte[] bytes, final String key, final int timeoutInMillis) throws Exception {
+ delegateBuffer.writeBytes(bytes, key, timeoutInMillis);
+ }
+
+ @Override
+ public Map.Entry, CheckpointState> read(final int timeoutInMillis) {
+ return delegateBuffer.read(timeoutInMillis);
+ }
+
+ @Override
+ public void checkpoint(final CheckpointState checkpointState) {
+ delegateBuffer.checkpoint(checkpointState);
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return delegateBuffer.isEmpty();
+ }
+
+ @Override
+ public boolean isByteBuffer() {
+ return delegateBuffer.isByteBuffer();
+ }
+
+ @Override
+ public Duration getDrainTimeout() {
+ return delegateBuffer.getDrainTimeout();
+ }
+
+ @Override
+ public boolean isWrittenOffHeapOnly() {
+ return delegateBuffer.isWrittenOffHeapOnly();
+ }
+
+ @Override
+ public void shutdown() {
+ delegateBuffer.shutdown();
+ }
+}
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java
index 6924cb35bc..a8ea4a3ee1 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java
@@ -161,6 +161,7 @@ public List getTypedList(final String attribute, final Class type) {
* Returns the value of the specified {@literal List>}, or {@code defaultValue} if this settings contains no value for
* the attribute.
*
+ * @param attribute attribute to be looked up
* @param keyType key type of the Map
* @param valueType value type stored in the Map
* @param The key type
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java
new file mode 100644
index 0000000000..743309bf75
--- /dev/null
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.event;
+
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import java.lang.ref.WeakReference;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.BiConsumer;
+import java.time.Instant;
+import java.io.Serializable;
+
+public class DefaultEventHandle implements EventHandle, InternalEventHandle, Serializable {
+ private Instant externalOriginationTime;
+ private final Instant internalOriginationTime;
+ private WeakReference acknowledgementSetRef;
+ private List> releaseConsumers;
+
+ public DefaultEventHandle(final Instant internalOriginationTime) {
+ this.acknowledgementSetRef = null;
+ this.externalOriginationTime = null;
+ this.internalOriginationTime = internalOriginationTime;
+ this.releaseConsumers = new ArrayList<>();
+ }
+
+ @Override
+ public void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet) {
+ this.acknowledgementSetRef = new WeakReference<>(acknowledgementSet);
+ }
+
+ @Override
+ public void setExternalOriginationTime(final Instant externalOriginationTime) {
+ this.externalOriginationTime = externalOriginationTime;
+ }
+
+ public AcknowledgementSet getAcknowledgementSet() {
+ if (acknowledgementSetRef == null) {
+ return null;
+ }
+ return acknowledgementSetRef.get();
+ }
+
+ @Override
+ public Instant getInternalOriginationTime() {
+ return this.internalOriginationTime;
+ }
+
+ @Override
+ public Instant getExternalOriginationTime() {
+ return this.externalOriginationTime;
+ }
+
+ @Override
+ public void release(boolean result) {
+ synchronized (releaseConsumers) {
+ for (final BiConsumer consumer: releaseConsumers) {
+ consumer.accept(this, result);
+ }
+ }
+ AcknowledgementSet acknowledgementSet = getAcknowledgementSet();
+ if (acknowledgementSet != null) {
+ acknowledgementSet.release(this, result);
+ }
+ }
+
+ @Override
+ public void onRelease(BiConsumer releaseConsumer) {
+ synchronized (releaseConsumers) {
+ releaseConsumers.add(releaseConsumer);
+ }
+ }
+}
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java
index e2ce55caa2..883297d567 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java
@@ -27,6 +27,8 @@ public class DefaultEventMetadata implements EventMetadata {
private final Instant timeReceived;
+ private Instant externalOriginationTime;
+
private Map attributes;
private Set tags;
@@ -43,6 +45,7 @@ private DefaultEventMetadata(final Builder builder) {
this.attributes = builder.attributes == null ? new HashMap<>() : new HashMap<>(builder.attributes);
this.tags = builder.tags == null ? new HashSet<>() : new HashSet(builder.tags);
+ this.externalOriginationTime = null;
}
private DefaultEventMetadata(final EventMetadata eventMetadata) {
@@ -50,6 +53,7 @@ private DefaultEventMetadata(final EventMetadata eventMetadata) {
this.timeReceived = eventMetadata.getTimeReceived();
this.attributes = new HashMap<>(eventMetadata.getAttributes());
this.tags = new HashSet<>(eventMetadata.getTags());
+ this.externalOriginationTime = null;
}
@Override
@@ -62,6 +66,16 @@ public Instant getTimeReceived() {
return timeReceived;
}
+ @Override
+ public Instant getExternalOriginationTime() {
+ return externalOriginationTime;
+ }
+
+ @Override
+ public void setExternalOriginationTime(Instant externalOriginationTime) {
+ this.externalOriginationTime = externalOriginationTime;
+ }
+
@Override
public Map getAttributes() {
return attributes;
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java
index 26f2d29e5f..d05dd8e36c 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java
@@ -5,6 +5,9 @@
package org.opensearch.dataprepper.model.event;
+import java.time.Instant;
+import java.util.function.BiConsumer;
+
public interface EventHandle {
/**
* releases event handle
@@ -14,4 +17,36 @@ public interface EventHandle {
* @since 2.2
*/
void release(boolean result);
+
+ /**
+ * sets external origination time
+ *
+ * @param externalOriginationTime externalOriginationTime to be set in the event handle
+ * @since 2.6
+ */
+ void setExternalOriginationTime(final Instant externalOriginationTime);
+
+ /**
+ * gets external origination time
+ *
+ * @return returns externalOriginationTime from the event handle. This can be null if it is never set.
+ * @since 2.6
+ */
+ Instant getExternalOriginationTime();
+
+ /**
+ * gets internal origination time
+ *
+ * @return returns internalOriginationTime from the event handle.
+ * @since 2.6
+ */
+ Instant getInternalOriginationTime();
+
+ /**
+ * registers onRelease consumer with event handle
+ *
+ * @param releaseConsumer consumer to be calledback when event handle is released.
+ */
+ void onRelease(BiConsumer releaseConsumer);
+
}
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java
index 511a87c1fa..5db56ba85c 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java
@@ -31,6 +31,20 @@ public interface EventMetadata extends Serializable {
*/
Instant getTimeReceived();
+ /**
+ * Returns the external origination time of the event
+ * @return the external origination time
+ * @since 2.6
+ */
+ Instant getExternalOriginationTime();
+
+ /**
+ * Sets the external origination time of the event
+ * @param externalOriginationTime the external origination time
+ * @since 2.6
+ */
+ void setExternalOriginationTime(Instant externalOriginationTime);
+
/**
* Returns the attributes
* @return a map of attributes
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java
new file mode 100644
index 0000000000..3817365f17
--- /dev/null
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.event;
+
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+
+public interface InternalEventHandle {
+ /**
+ * sets acknowledgement set
+ *
+ * @param acknowledgementSet acknowledgementSet to be set in the event handle
+ * @since 2.6
+ */
+ void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet);
+
+ /**
+ * gets acknowledgement set
+ *
+ * @return returns acknowledgementSet from the event handle
+ * @since 2.6
+ */
+ AcknowledgementSet getAcknowledgementSet();
+
+}
+
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
index 686ee1d59f..341d5277bc 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
@@ -91,11 +91,13 @@ protected JacksonEvent(final Builder builder) {
}
this.jsonNode = getInitialJsonNode(builder.data);
+ this.eventHandle = new DefaultEventHandle(eventMetadata.getTimeReceived());
}
protected JacksonEvent(final JacksonEvent otherEvent) {
this.jsonNode = otherEvent.jsonNode.deepCopy();
this.eventMetadata = DefaultEventMetadata.fromEventMetadata(otherEvent.eventMetadata);
+ this.eventHandle = new DefaultEventHandle(eventMetadata.getTimeReceived());
}
public static Event fromMessage(String message) {
@@ -152,10 +154,6 @@ public void put(final String key, final Object value) {
}
}
- public void setEventHandle(EventHandle handle) {
- this.eventHandle = handle;
- }
-
@Override
public EventHandle getEventHandle() {
return eventHandle;
@@ -320,28 +318,6 @@ public String formatString(final String format, final ExpressionEvaluator expres
return formatStringInternal(format, expressionEvaluator);
}
- public static boolean isValidFormatExpressions(final String format, final ExpressionEvaluator expressionEvaluator) {
- if (Objects.isNull(expressionEvaluator)) {
- return false;
- }
- int fromIndex = 0;
- int position = 0;
- while ((position = format.indexOf("${", fromIndex)) != -1) {
- int endPosition = format.indexOf("}", position + 1);
- if (endPosition == -1) {
- return false;
- }
- String name = format.substring(position + 2, endPosition);
-
- Object val;
- if (!expressionEvaluator.isValidExpressionStatement(name)) {
- return false;
- }
- fromIndex = endPosition + 1;
- }
- return true;
- }
-
private String formatStringInternal(final String format, final ExpressionEvaluator expressionEvaluator) {
int fromIndex = 0;
String result = "";
@@ -363,7 +339,7 @@ private String formatStringInternal(final String format, final ExpressionEvaluat
}
if (val == null) {
- if (Objects.nonNull(expressionEvaluator) && expressionEvaluator.isValidExpressionStatement(name)) {
+ if (expressionEvaluator != null && expressionEvaluator.isValidExpressionStatement(name)) {
val = expressionEvaluator.evaluate(name, this);
} else {
throw new EventKeyNotFoundException(String.format("The key %s could not be found in the Event when formatting", name));
@@ -411,12 +387,21 @@ public Map toMap() {
return mapper.convertValue(jsonNode, MAP_TYPE_REFERENCE);
}
+
+ public static boolean isValidEventKey(final String key) {
+ try {
+ checkKey(key);
+ return true;
+ } catch (final Exception e) {
+ return false;
+ }
+ }
private String checkAndTrimKey(final String key) {
checkKey(key);
return trimKey(key);
}
- private void checkKey(final String key) {
+ private static void checkKey(final String key) {
checkNotNull(key, "key cannot be null");
checkArgument(!key.isEmpty(), "key cannot be an empty string");
if (key.length() > MAX_KEY_LENGTH) {
@@ -433,7 +418,7 @@ private String trimKey(final String key) {
return trimmedLeadingSlash.endsWith(SEPARATOR) ? trimmedLeadingSlash.substring(0, trimmedLeadingSlash.length() - 2) : trimmedLeadingSlash;
}
- private boolean isValidKey(final String key) {
+ private static boolean isValidKey(final String key) {
for (int i = 0; i < key.length(); i++) {
char c = key.charAt(i);
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java
index b52c850e46..b865ce0eb5 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java
@@ -27,8 +27,8 @@ public class JacksonExponentialHistogram extends JacksonMetric implements Expone
private static final String SCALE_KEY = "scale";
private static final String AGGREGATION_TEMPORALITY_KEY = "aggregationTemporality";
private static final String ZERO_COUNT_KEY = "zeroCount";
- private static final String POSITIVE_BUCKETS_KEY = "positiveBuckets";
- private static final String NEGATIVE_BUCKETS_KEY = "negativeBuckets";
+ public static final String POSITIVE_BUCKETS_KEY = "positiveBuckets";
+ public static final String NEGATIVE_BUCKETS_KEY = "negativeBuckets";
private static final String NEGATIVE_KEY = "negative";
private static final String POSITIVE_KEY = "positive";
private static final String NEGATIVE_OFFSET_KEY = "negativeOffset";
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java
index 0209f7012d..f9e066875d 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java
@@ -29,7 +29,7 @@ public class JacksonHistogram extends JacksonMetric implements Histogram {
private static final String AGGREGATION_TEMPORALITY_KEY = "aggregationTemporality";
private static final String BUCKET_COUNTS_KEY = "bucketCounts";
private static final String EXPLICIT_BOUNDS_COUNT_KEY = "explicitBoundsCount";
- private static final String BUCKETS_KEY = "buckets";
+ public static final String BUCKETS_KEY = "buckets";
private static final String BUCKET_COUNTS_LIST_KEY = "bucketCountsList";
private static final String EXPLICIT_BOUNDS_KEY = "explicitBounds";
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java
index 0ab81ed7e0..8d8ebf0f87 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java
@@ -28,7 +28,7 @@ public abstract class JacksonMetric extends JacksonEvent implements Metric {
protected static final String SERVICE_NAME_KEY = "serviceName";
protected static final String KIND_KEY = "kind";
protected static final String UNIT_KEY = "unit";
- protected static final String ATTRIBUTES_KEY = "attributes";
+ public static final String ATTRIBUTES_KEY = "attributes";
protected static final String SCHEMA_URL_KEY = "schemaUrl";
protected static final String EXEMPLARS_KEY = "exemplars";
protected static final String FLAGS_KEY = "flags";
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java
index 6a62e2647a..3c436fc68a 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java
@@ -10,6 +10,9 @@ public interface PluginConfigObservable {
/**
* Onboard a new {@link PluginConfigObserver} within the plugin.
+ *
+ * @param pluginConfigObserver plugin config observer
+ * @return returns true if the opration is successful, false otherwise
*/
boolean addPluginConfigObserver(PluginConfigObserver pluginConfigObserver);
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java
index 66982a8db6..ae3ff0ccb4 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java
@@ -7,6 +7,9 @@
public interface PluginConfigPublisher {
/**
* Onboard a new {@link PluginConfigObservable}.
+ *
+ * @param pluginConfigObservable observable plugin configuration
+ * @return true if the operation is successful, false otherwise
*/
boolean addPluginConfigObservable(PluginConfigObservable pluginConfigObservable);
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java
index b233d89043..155df776f4 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java
@@ -22,6 +22,7 @@ public interface PluginFactory {
*
* @param baseClass The class type that the plugin is supporting.
* @param pluginSetting The {@link PluginSetting} to configure this plugin
+ * @param args variable number of arguments
* @param The type
* @return A new instance of your plugin, configured
* @since 1.2
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java
index a5bb6e8704..65ae2d6ce5 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java
@@ -12,7 +12,6 @@
* TODO: The current implementation focuses on proving the bare bones for which this class only need to
* TODO: support sample test cases.
*/
-@Deprecated
public class Record {
private final T data;
private final RecordMetadata metadata;
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java
index 823abd30e1..8efbe0960e 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java
@@ -13,7 +13,6 @@
* The RecordMetadata class provides a wrapper around the ImmutableMap making metadata management easier for the
* user to access.
*/
-@Deprecated
public class RecordMetadata {
private static final RecordMetadata DEFAULT_METADATA = new RecordMetadata();
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java
index f116d89e9e..1c3e596265 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java
@@ -11,6 +11,7 @@
import org.opensearch.dataprepper.metrics.PluginMetrics;
import org.opensearch.dataprepper.model.configuration.PluginSetting;
import org.opensearch.dataprepper.model.record.Record;
+import org.opensearch.dataprepper.model.event.Event;
import java.util.Collection;
@@ -22,6 +23,7 @@ public abstract class AbstractSink> implements Sink {
protected static final int DEFAULT_WAIT_TIME_MS = 1000;
protected final PluginMetrics pluginMetrics;
private final Counter recordsInCounter;
+ private final SinkLatencyMetrics latencyMetrics;
private final Timer timeElapsedTimer;
private Thread retryThread;
private int maxRetries;
@@ -31,6 +33,7 @@ public AbstractSink(final PluginSetting pluginSetting, int numRetries, int waitT
this.pluginMetrics = PluginMetrics.fromPluginSetting(pluginSetting);
recordsInCounter = pluginMetrics.counter(MetricNames.RECORDS_IN);
timeElapsedTimer = pluginMetrics.timer(MetricNames.TIME_ELAPSED);
+ this.latencyMetrics = new SinkLatencyMetrics(pluginMetrics);
retryThread = null;
this.maxRetries = numRetries;
this.waitTimeMs = waitTimeMs;
@@ -77,6 +80,20 @@ public void shutdown() {
}
}
+ @Override
+ public void updateLatencyMetrics(Collection records) {
+ for (final Record record : records) {
+ if (record.getData() instanceof Event) {
+ Event event = (Event)record.getData();
+ event.getEventHandle().onRelease((eventHandle, result) -> {
+ if (result) {
+ latencyMetrics.update(eventHandle);
+ }
+ });
+ }
+ }
+ }
+
Thread.State getRetryThreadState() {
if (retryThread != null) {
return retryThread.getState();
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java
index 0ce6fa5ac1..178566ba5b 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java
@@ -38,4 +38,12 @@ public interface Sink> {
*/
boolean isReady();
+ /**
+ * updates latency metrics of sink
+ *
+ * @param events list of events used for updating the latency metrics
+ */
+ default void updateLatencyMetrics(final Collection events) {
+ }
+
}
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java
new file mode 100644
index 0000000000..3a39c75b96
--- /dev/null
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.sink;
+
+import org.opensearch.dataprepper.metrics.PluginMetrics;
+import io.micrometer.core.instrument.DistributionSummary;
+import org.opensearch.dataprepper.model.event.EventHandle;
+
+import java.time.Duration;
+import java.time.Instant;
+
+public class SinkLatencyMetrics {
+ public static final String INTERNAL_LATENCY = "PipelineLatency";
+ public static final String EXTERNAL_LATENCY = "EndToEndLatency";
+ private final DistributionSummary internalLatencySummary;
+ private final DistributionSummary externalLatencySummary;
+
+ public SinkLatencyMetrics(PluginMetrics pluginMetrics) {
+ internalLatencySummary = pluginMetrics.summary(INTERNAL_LATENCY);
+ externalLatencySummary = pluginMetrics.summary(EXTERNAL_LATENCY);
+ }
+ public void update(final EventHandle eventHandle) {
+ Instant now = Instant.now();
+ internalLatencySummary.record(Duration.between(eventHandle.getInternalOriginationTime(), now).toMillis());
+ if (eventHandle.getExternalOriginationTime() == null) {
+ return;
+ }
+ externalLatencySummary.record(Duration.between(eventHandle.getExternalOriginationTime(), now).toMillis());
+ }
+}
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java
index 9341a13e59..5423732b6f 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java
@@ -46,9 +46,10 @@ public interface EnhancedSourceCoordinator {
*
* @param partition The partition to be updated.
* @param The progress state class
+ * @param ownershipTimeoutRenewal The amount of time to update ownership of the partition before another instance can acquire it.
* @throws org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException when the partition was not updated successfully
*/
- void saveProgressStateForPartition(EnhancedSourcePartition partition);
+ void saveProgressStateForPartition(EnhancedSourcePartition partition, Duration ownershipTimeoutRenewal);
/**
* This method is used to release the lease of a partition in the coordination store.
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java
index ac52ace478..0fa683b2a5 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java
@@ -48,6 +48,11 @@ public void setSourcePartitionStoreItem(SourcePartitionStoreItem sourcePartition
/**
* Helper method to convert progress state.
* This is because the state is currently stored as a String in the coordination store.
+ *
+ * @param progressStateClass class of progress state
+ * @param serializedPartitionProgressState serialized value of the partition progress state
+ *
+ * @return returns the converted value of the progress state
*/
public T convertStringToPartitionProgressState(Class progressStateClass, final String serializedPartitionProgressState) {
if (Objects.isNull(serializedPartitionProgressState)) {
@@ -69,6 +74,9 @@ public T convertStringToPartitionProgressState(Class progressStateClass, fina
/**
* Helper method to convert progress state to String
* This is because the state is currently stored as a String in the coordination store.
+ *
+ * @param partitionProgressState optional parameter indicating the partition progress state
+ * @return returns the progress state as string
*/
public String convertPartitionProgressStatetoString(Optional partitionProgressState) {
if (partitionProgressState.isEmpty()) {
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java
index ae09e953ff..1681563675 100644
--- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java
+++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java
@@ -21,6 +21,7 @@
*/
public class ByteCount {
private static final Pattern BYTE_PATTERN = Pattern.compile("^(?\\d+\\.?\\d*)(?[a-z]+)?\\z");
+ private static final ByteCount ZERO_BYTES = new ByteCount(0);
private final long bytes;
private ByteCount(final long bytes) {
@@ -94,6 +95,10 @@ public static ByteCount parse(final String string) {
return new ByteCount(byteCount.longValue());
}
+ public static ByteCount zeroBytes() {
+ return ZERO_BYTES;
+ }
+
private static BigDecimal scaleToBytes(final BigDecimal value, final Unit unit) {
return value.multiply(BigDecimal.valueOf(unit.multiplier));
}
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java
index 36a60ac447..9b76fbc807 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java
@@ -23,6 +23,11 @@ public Object evaluate(final String statement, final Event event) {
public Boolean isValidExpressionStatement(final String statement) {
return true;
}
+
+ @Override
+ public Boolean isValidFormatExpression(String format) {
+ return true;
+ }
}
@Test
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java
index bc9df29ba3..505a00cc4a 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java
@@ -5,43 +5,55 @@
package org.opensearch.dataprepper.model.buffer;
-import org.junit.Assert;
import org.junit.jupiter.api.Test;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.record.Record;
import java.time.Duration;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.spy;
-public class BufferTest {
+class BufferTest {
+
+ private Buffer createObjectUnderTest() {
+ return spy(Buffer.class);
+ }
@Test
- public void testGetDrainTimeout() {
- final Buffer> buffer = spy(Buffer.class);
+ void testGetDrainTimeout() {
+ final Buffer> buffer = createObjectUnderTest();
- Assert.assertEquals(Duration.ZERO, buffer.getDrainTimeout());
+ assertEquals(Duration.ZERO, buffer.getDrainTimeout());
}
@Test
- public void testShutdown() {
- final Buffer> buffer = spy(Buffer.class);
+ void testShutdown() {
+ final Buffer> buffer = createObjectUnderTest();
buffer.shutdown();
}
@Test
- public void testIsByteBuffer() {
- final Buffer> buffer = spy(Buffer.class);
+ void testIsByteBuffer() {
+ final Buffer> buffer = createObjectUnderTest();
- Assert.assertEquals(false, buffer.isByteBuffer());
+ assertEquals(false, buffer.isByteBuffer());
+ }
+
+ @Test
+ void isWrittenOffHeapOnly_returns_false_by_default() {
+ assertThat(createObjectUnderTest().isWrittenOffHeapOnly(), equalTo(false));
}
@Test
- public void testWriteBytes() {
- final Buffer> buffer = spy(Buffer.class);
+ void testWriteBytes() {
+ final Buffer> buffer = createObjectUnderTest();
byte[] bytes = new byte[2];
- Assert.assertThrows(RuntimeException.class, () -> buffer.writeBytes(bytes, "", 10));
+ assertThrows(UnsupportedOperationException.class, () -> buffer.writeBytes(bytes, "", 10));
}
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java
new file mode 100644
index 0000000000..09e300c206
--- /dev/null
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java
@@ -0,0 +1,211 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.buffer;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.dataprepper.model.CheckpointState;
+import org.opensearch.dataprepper.model.record.Record;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.UUID;
+import java.util.concurrent.TimeoutException;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.sameInstance;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class DelegatingBufferTest {
+ @Mock
+ private Buffer> innerBuffer;
+
+ @Mock
+ private Record record;
+
+ private Collection> records;
+
+ private int timeoutInMillis;
+ private Random random;
+
+ @BeforeEach
+ void setUp() {
+ random = new Random();
+ timeoutInMillis = random.nextInt(10_000) + 100;
+
+ records = List.of(record, mock(Record.class));
+ }
+
+ private static class TestDelegatingBuffer extends DelegatingBuffer> {
+ TestDelegatingBuffer(final Buffer> delegateBuffer) {
+ super(delegateBuffer);
+ }
+ }
+
+ private DelegatingBuffer> createObjectUnderTest() {
+ return new TestDelegatingBuffer(innerBuffer);
+ }
+
+ @Test
+ void constructor_throws_if_delegate_is_null() {
+ innerBuffer = null;
+
+ assertThrows(NullPointerException.class, this::createObjectUnderTest);
+ }
+
+ @Test
+ void write_calls_inner_write() throws TimeoutException {
+ createObjectUnderTest().write(record, timeoutInMillis);
+
+ verify(innerBuffer).write(record, timeoutInMillis);
+ }
+
+ @ParameterizedTest
+ @ValueSource(classes = {RuntimeException.class, TimeoutException.class})
+ void write_throws_exceptions_from_inner_write(final Class exceptionType) throws TimeoutException {
+ final Throwable exception = mock(exceptionType);
+ doThrow(exception).when(innerBuffer).write(any(), anyInt());
+
+ final DelegatingBuffer> objectUnderTest = createObjectUnderTest();
+ final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.write(record, timeoutInMillis));
+
+ assertThat(actualException, sameInstance(exception));
+ }
+
+ @Test
+ void writeAll_calls_inner_writeAll() throws Exception {
+ createObjectUnderTest().writeAll(records, timeoutInMillis);
+
+ verify(innerBuffer).writeAll(records, timeoutInMillis);
+ }
+
+ @ParameterizedTest
+ @ValueSource(classes = {Exception.class, RuntimeException.class, TimeoutException.class})
+ void writeAll_throws_exceptions_from_inner_writeAll(final Class exceptionType) throws Exception {
+ final Throwable exception = mock(exceptionType);
+ doThrow(exception).when(innerBuffer).writeAll(any(), anyInt());
+
+ final DelegatingBuffer> objectUnderTest = createObjectUnderTest();
+ final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.writeAll(records, timeoutInMillis));
+
+ assertThat(actualException, sameInstance(exception));
+ }
+
+ @Test
+ void writeBytes_calls_inner_writeBytes() throws Exception {
+ final byte[] bytesToWrite = new byte[64];
+ random.nextBytes(bytesToWrite);
+ final String key = UUID.randomUUID().toString();
+ createObjectUnderTest().writeBytes(bytesToWrite, key, timeoutInMillis);
+
+ verify(innerBuffer).writeBytes(bytesToWrite, key, timeoutInMillis);
+ }
+
+ @ParameterizedTest
+ @ValueSource(classes = {Exception.class, RuntimeException.class, TimeoutException.class})
+ void writeBytes_throws_exceptions_from_inner_writeBytes(final Class exceptionType) throws Exception {
+ final Throwable exception = mock(exceptionType);
+ doThrow(exception).when(innerBuffer).writeBytes(any(), any(), anyInt());
+
+ final byte[] bytesToWrite = new byte[64];
+ random.nextBytes(bytesToWrite);
+ final String key = UUID.randomUUID().toString();
+
+ final DelegatingBuffer> objectUnderTest = createObjectUnderTest();
+ final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.writeBytes(bytesToWrite, key, timeoutInMillis));
+
+ assertThat(actualException, sameInstance(exception));
+ }
+
+ @Test
+ void read_returns_inner_read() {
+ final Map.Entry>, CheckpointState> readResult = mock(Map.Entry.class);
+ when(innerBuffer.read(timeoutInMillis)).thenReturn(readResult);
+
+ assertThat(createObjectUnderTest().read(timeoutInMillis),
+ equalTo(readResult));
+ }
+
+ @Test
+ void read_throws_exceptions_from_inner_read() {
+ final RuntimeException exception = mock(RuntimeException.class);
+
+ when(innerBuffer.read(timeoutInMillis)).thenThrow(exception);
+
+ final DelegatingBuffer> objectUnderTest = createObjectUnderTest();
+ final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.read(timeoutInMillis));
+
+ assertThat(actualException, sameInstance(exception));
+ }
+
+ @Test
+ void checkpoint_calls_inner_checkpoint() {
+ final CheckpointState checkpointState = mock(CheckpointState.class);
+
+ createObjectUnderTest().checkpoint(checkpointState);
+
+ verify(innerBuffer).checkpoint(checkpointState);
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ void isEmpty_returns_inner_isEmpty(final boolean isEmpty) {
+ when(innerBuffer.isEmpty()).thenReturn(isEmpty);
+
+ assertThat(createObjectUnderTest().isEmpty(),
+ equalTo(isEmpty));
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ void isByteBuffer_returns_inner_isByteBuffer(final boolean isByteBuffer) {
+ when(innerBuffer.isByteBuffer()).thenReturn(isByteBuffer);
+
+ assertThat(createObjectUnderTest().isByteBuffer(),
+ equalTo(isByteBuffer));
+ }
+
+ @Test
+ void getDrainTimeout_returns_inner_getDrainTimeout() {
+ final Duration drainTimeout = Duration.ofSeconds(random.nextInt(10_000) + 100);
+ when(innerBuffer.getDrainTimeout()).thenReturn(drainTimeout);
+
+ assertThat(createObjectUnderTest().getDrainTimeout(),
+ equalTo(drainTimeout));
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ void isWrittenOffHeapOnly_returns_inner_isWrittenOffHeapOnly(final boolean isWrittenOffHeapOnly) {
+ when(innerBuffer.isWrittenOffHeapOnly()).thenReturn(isWrittenOffHeapOnly);
+
+ assertThat(createObjectUnderTest().isWrittenOffHeapOnly(),
+ equalTo(isWrittenOffHeapOnly));
+ }
+
+ @Test
+ void shutdown_calls_inner_shutdown() {
+ createObjectUnderTest().shutdown();
+
+ verify(innerBuffer).shutdown();
+ }
+}
\ No newline at end of file
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java
new file mode 100644
index 0000000000..b2a66b2d1d
--- /dev/null
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.event;
+
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import org.junit.jupiter.api.Test;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verify;
+import static org.mockito.ArgumentMatchers.any;
+import org.mockito.Mock;
+
+import java.time.Instant;
+
+class DefaultEventHandleTests {
+ @Mock
+ private AcknowledgementSet acknowledgementSet;
+ private int count;
+
+ @Test
+ void testBasic() {
+ Instant now = Instant.now();
+ DefaultEventHandle eventHandle = new DefaultEventHandle(now);
+ assertThat(eventHandle.getAcknowledgementSet(), equalTo(null));
+ assertThat(eventHandle.getInternalOriginationTime(), equalTo(now));
+ assertThat(eventHandle.getExternalOriginationTime(), equalTo(null));
+ eventHandle.release(true);
+ }
+
+ @Test
+ void testWithAcknowledgementSet() {
+ acknowledgementSet = mock(AcknowledgementSet.class);
+ when(acknowledgementSet.release(any(EventHandle.class), any(Boolean.class))).thenReturn(true);
+ Instant now = Instant.now();
+ DefaultEventHandle eventHandle = new DefaultEventHandle(now);
+ assertThat(eventHandle.getAcknowledgementSet(), equalTo(null));
+ assertThat(eventHandle.getInternalOriginationTime(), equalTo(now));
+ assertThat(eventHandle.getExternalOriginationTime(), equalTo(null));
+ eventHandle.setAcknowledgementSet(acknowledgementSet);
+ eventHandle.release(true);
+ verify(acknowledgementSet).release(eventHandle, true);
+ }
+
+ @Test
+ void testWithExternalOriginationTime() {
+ Instant now = Instant.now();
+ DefaultEventHandle eventHandle = new DefaultEventHandle(now);
+ assertThat(eventHandle.getAcknowledgementSet(), equalTo(null));
+ assertThat(eventHandle.getInternalOriginationTime(), equalTo(now));
+ assertThat(eventHandle.getExternalOriginationTime(), equalTo(null));
+ eventHandle.setExternalOriginationTime(now.minusSeconds(60));
+ assertThat(eventHandle.getExternalOriginationTime(), equalTo(now.minusSeconds(60)));
+ eventHandle.release(true);
+ }
+
+ @Test
+ void testWithOnReleaseHandler() {
+ Instant now = Instant.now();
+ count = 0;
+ DefaultEventHandle eventHandle = new DefaultEventHandle(now);
+ eventHandle.onRelease((handle, result) -> {if (result) count++; });
+ eventHandle.release(true);
+ assertThat(count, equalTo(1));
+
+ }
+
+}
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java
index fa624a9e2b..479e7be0c2 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java
@@ -24,6 +24,7 @@
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.not;
import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
import static org.hamcrest.CoreMatchers.sameInstance;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.anEmptyMap;
@@ -80,6 +81,16 @@ public void testGetTimeReceived() {
assertThat(timeReceived, is(equalTo(testTimeReceived)));
}
+ @Test
+ public void testExternalOriginationTime() {
+ Instant externalOriginationTime = eventMetadata.getExternalOriginationTime();
+ assertThat(externalOriginationTime, is(nullValue()));
+ Instant now = Instant.now();
+ eventMetadata.setExternalOriginationTime(now);
+ externalOriginationTime = eventMetadata.getExternalOriginationTime();
+ assertThat(externalOriginationTime, is(equalTo(now)));
+ }
+
@Test
public void testGetAttributes() {
final Map attributes = eventMetadata.getAttributes();
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java
index 62ce3dc48d..4fe8f272cc 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java
@@ -42,12 +42,6 @@
public class JacksonEventTest {
- class TestEventHandle implements EventHandle {
- @Override
- public void release(boolean result) {
- }
- }
-
private Event event;
private String eventType;
@@ -398,6 +392,8 @@ public void testBuild_withEventType() {
.build();
assertThat(event.getMetadata().getEventType(), is(equalTo(eventType)));
+ assertThat(event.getEventHandle(), is(notNullValue()));
+ assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue()));
}
@Test
@@ -411,6 +407,8 @@ public void testBuild_withTimeReceived() {
.build();
assertThat(event.getMetadata().getTimeReceived(), is(equalTo(now)));
+ assertThat(event.getEventHandle(), is(notNullValue()));
+ assertThat(event.getEventHandle().getInternalOriginationTime(), is(equalTo(now)));
}
@Test
@@ -422,6 +420,8 @@ public void testBuild_withMessageValue() {
assertThat(event, is(notNullValue()));
assertThat(event.get("message", String.class), is(equalTo(message)));
+ assertThat(event.getEventHandle(), is(notNullValue()));
+ assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue()));
}
@Test
@@ -535,25 +535,6 @@ public void testBuild_withFormatString(String formattedString, String finalStrin
assertThat(event.formatString(formattedString), is(equalTo(finalString)));
}
- @ParameterizedTest
- @CsvSource({
- "abc-${/foo, false",
- "abc-${/foo}, true",
- "abc-${getMetadata(\"key\")}, true",
- "abc-${getXYZ(\"key\")}, false"
- })
- public void testBuild_withIsValidFormatExpressions(final String format, final Boolean expectedResult) {
- final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class);
- when(expressionEvaluator.isValidExpressionStatement("/foo")).thenReturn(true);
- when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"key\")")).thenReturn(true);
- assertThat(JacksonEvent.isValidFormatExpressions(format, expressionEvaluator), equalTo(expectedResult));
- }
-
- @Test
- public void testBuild_withIsValidFormatExpressionsWithNullEvaluator() {
- assertThat(JacksonEvent.isValidFormatExpressions("${}", null), equalTo(false));
- }
-
@Test
public void formatString_with_expression_evaluator_catches_exception_when_Event_get_throws_exception() {
@@ -678,6 +659,8 @@ void fromEvent_with_a_JacksonEvent() {
assertThat(createdEvent, notNullValue());
assertThat(createdEvent, not(sameInstance(originalEvent)));
+ assertThat(event.getEventHandle(), is(notNullValue()));
+ assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue()));
assertThat(createdEvent.toMap(), equalTo(dataObject));
assertThat(createdEvent.getJsonNode(), not(sameInstance(originalEvent.getJsonNode())));
@@ -707,19 +690,6 @@ void fromEvent_with_a_non_JacksonEvent() {
assertThat(createdEvent.getMetadata(), equalTo(eventMetadata));
}
- @Test
- void testEventHandleGetAndSet() {
- EventHandle testEventHandle = new TestEventHandle();
- final String jsonString = "{\"foo\": \"bar\"}";
-
- final JacksonEvent event = JacksonEvent.builder()
- .withEventType(eventType)
- .withData(jsonString)
- .build();
- event.setEventHandle(testEventHandle);
- assertThat(event.getEventHandle(), equalTo(testEventHandle));
- }
-
@Test
void testJsonStringBuilder() {
final String jsonString = "{\"foo\":\"bar\"}";
@@ -829,6 +799,11 @@ void testJsonStringBuilderWithExcludeKeys() {
}
+ @ParameterizedTest
+ @CsvSource(value = {"test_key, true", "/test_key, true", "inv(alid, false", "getMetadata(\"test_key\"), false"})
+ void isValidEventKey_returns_expected_result(final String key, final boolean isValid) {
+ assertThat(JacksonEvent.isValidEventKey(key), equalTo(isValid));
+ }
private static Map createComplexDataMap() {
final Map dataObject = new HashMap<>();
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java
index e4b19cf7ca..6d58f7dd71 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java
@@ -13,6 +13,12 @@
import org.opensearch.dataprepper.metrics.MetricsTestUtil;
import org.opensearch.dataprepper.model.configuration.PluginSetting;
import org.opensearch.dataprepper.model.record.Record;
+import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.event.JacksonEvent;
+import org.opensearch.dataprepper.model.event.EventHandle;
+
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.mock;
import java.time.Duration;
import java.util.Arrays;
@@ -25,6 +31,7 @@
import static org.awaitility.Awaitility.await;
public class AbstractSinkTest {
+ private int count;
@Test
public void testMetrics() {
final String sinkName = "testSink";
@@ -35,6 +42,8 @@ public void testMetrics() {
AbstractSink> abstractSink = new AbstractSinkImpl(pluginSetting);
abstractSink.initialize();
Assert.assertEquals(abstractSink.isReady(), true);
+ abstractSink.updateLatencyMetrics(Arrays.asList(
+ new Record<>(UUID.randomUUID().toString())));
abstractSink.output(Arrays.asList(
new Record<>(UUID.randomUUID().toString()),
new Record<>(UUID.randomUUID().toString()),
@@ -80,6 +89,61 @@ public void testSinkNotReady() {
abstractSink.shutdown();
}
+ @Test
+ public void testSinkWithRegisterEventReleaseHandler() {
+ final String sinkName = "testSink";
+ final String pipelineName = "pipelineName";
+ MetricsTestUtil.initMetrics();
+ PluginSetting pluginSetting = new PluginSetting(sinkName, Collections.emptyMap());
+ pluginSetting.setPipelineName(pipelineName);
+ AbstractSink> abstractSink = new AbstractEventSinkImpl(pluginSetting);
+ abstractSink.initialize();
+ Assert.assertEquals(abstractSink.isReady(), true);
+ count = 0;
+ Event event = JacksonEvent.builder()
+ .withEventType("event")
+ .build();
+ Record record = mock(Record.class);
+ EventHandle eventHandle = mock(EventHandle.class);
+ when(record.getData()).thenReturn(event);
+
+ abstractSink.updateLatencyMetrics(Arrays.asList(record));
+ abstractSink.output(Arrays.asList(record));
+ await().atMost(Duration.ofSeconds(5))
+ .until(abstractSink::isReady);
+ abstractSink.shutdown();
+ }
+
+ private static class AbstractEventSinkImpl extends AbstractSink> {
+
+ public AbstractEventSinkImpl(PluginSetting pluginSetting) {
+ super(pluginSetting, 10, 1000);
+ }
+
+ @Override
+ public void doOutput(Collection> records) {
+ for (final Record record: records) {
+ Event event = record.getData();
+ event.getEventHandle().release(true);
+ }
+ }
+
+ @Override
+ public void shutdown() {
+ super.shutdown();
+ }
+
+ @Override
+ public void doInitialize() {
+ }
+
+ @Override
+ public boolean isReady() {
+ return true;
+ }
+ }
+
+
private static class AbstractSinkImpl extends AbstractSink> {
public AbstractSinkImpl(PluginSetting pluginSetting) {
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java
new file mode 100644
index 0000000000..4cf5043cae
--- /dev/null
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.sink;
+
+import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.event.EventHandle;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import io.micrometer.core.instrument.DistributionSummary;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+
+import java.time.Instant;
+
+class SinkLatencyMetricsTest {
+
+ private PluginMetrics pluginMetrics;
+ private EventHandle eventHandle;
+ private SinkLatencyMetrics latencyMetrics;
+ private DistributionSummary internalLatencySummary;
+ private DistributionSummary externalLatencySummary;
+
+ public SinkLatencyMetrics createObjectUnderTest() {
+ return new SinkLatencyMetrics(pluginMetrics);
+ }
+
+ @BeforeEach
+ void setup() {
+ pluginMetrics = mock(PluginMetrics.class);
+ SimpleMeterRegistry registry = new SimpleMeterRegistry();
+ internalLatencySummary = DistributionSummary
+ .builder("internalLatency")
+ .baseUnit("milliseconds")
+ .register(registry);
+ externalLatencySummary = DistributionSummary
+ .builder("externalLatency")
+ .baseUnit("milliseconds")
+ .register(registry);
+ when(pluginMetrics.summary(SinkLatencyMetrics.INTERNAL_LATENCY)).thenReturn(internalLatencySummary);
+ when(pluginMetrics.summary(SinkLatencyMetrics.EXTERNAL_LATENCY)).thenReturn(externalLatencySummary);
+ eventHandle = mock(EventHandle.class);
+ when(eventHandle.getInternalOriginationTime()).thenReturn(Instant.now());
+ latencyMetrics = createObjectUnderTest();
+ }
+
+ @Test
+ public void testInternalOriginationTime() {
+ latencyMetrics.update(eventHandle);
+ assertThat(internalLatencySummary.count(), equalTo(1L));
+ }
+
+ @Test
+ public void testExternalOriginationTime() {
+ when(eventHandle.getExternalOriginationTime()).thenReturn(Instant.now().minusMillis(10));
+ latencyMetrics.update(eventHandle);
+ assertThat(internalLatencySummary.count(), equalTo(1L));
+ assertThat(externalLatencySummary.count(), equalTo(1L));
+ assertThat(externalLatencySummary.max(), greaterThanOrEqualTo(10.0));
+ }
+}
+
+
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java
new file mode 100644
index 0000000000..5f66a623aa
--- /dev/null
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.sink;
+
+import org.opensearch.dataprepper.model.record.Record;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collection;
+import java.util.Collections;
+
+public class SinkTest {
+ private static class SinkTestClass implements Sink> {
+
+ @Override
+ public boolean isReady() {
+ return true;
+ }
+
+ @Override
+ public void shutdown() {
+ }
+
+ @Override
+ public void initialize() {
+ }
+
+ @Override
+ public void output(Collection> records) {
+ }
+
+ };
+
+ SinkTestClass sink;
+
+ @Test
+ public void testSinkUpdateLatencyMetrics() {
+ sink = new SinkTestClass();
+ sink.updateLatencyMetrics(Collections.emptyList());
+ }
+}
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java
index eda34eae69..b717289a7f 100644
--- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java
+++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java
@@ -5,6 +5,7 @@
package org.opensearch.dataprepper.model.types;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;
import org.junit.jupiter.params.provider.ValueSource;
@@ -13,6 +14,7 @@
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.sameInstance;
import static org.junit.jupiter.api.Assertions.assertThrows;
class ByteCountTest {
@@ -145,4 +147,16 @@ void parse_returns_rounded_bytes_for_implicit_fractional_bytes(final String byte
assertThat(byteCount, notNullValue());
assertThat(byteCount.getBytes(), equalTo(expectedBytes));
}
+
+ @Test
+ void zeroBytes_returns_bytes_with_getBytes_equal_to_0() {
+ assertThat(ByteCount.zeroBytes(), notNullValue());
+ assertThat(ByteCount.zeroBytes().getBytes(), equalTo(0L));
+ }
+
+ @Test
+ void zeroBytes_returns_same_instance() {
+ assertThat(ByteCount.zeroBytes(), notNullValue());
+ assertThat(ByteCount.zeroBytes(), sameInstance(ByteCount.zeroBytes()));
+ }
}
\ No newline at end of file
diff --git a/data-prepper-core/build.gradle b/data-prepper-core/build.gradle
index 72d9c47c54..1400d72ad2 100644
--- a/data-prepper-core/build.gradle
+++ b/data-prepper-core/build.gradle
@@ -27,7 +27,7 @@ dependencies {
implementation 'javax.ws.rs:javax.ws.rs-api:2.1.1'
implementation 'software.amazon.awssdk:cloudwatch'
implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final'
- implementation platform('org.apache.logging.log4j:log4j-bom:2.20.0')
+ implementation platform('org.apache.logging.log4j:log4j-bom:2.21.1')
implementation 'org.apache.logging.log4j:log4j-core'
implementation 'org.apache.logging.log4j:log4j-slf4j2-impl'
implementation 'javax.inject:javax.inject:1'
@@ -39,7 +39,7 @@ dependencies {
exclude group: 'commons-logging', module: 'commons-logging'
}
implementation 'software.amazon.cloudwatchlogs:aws-embedded-metrics:2.0.0-beta-1'
- testImplementation 'org.apache.logging.log4j:log4j-jpl:2.20.0'
+ testImplementation 'org.apache.logging.log4j:log4j-jpl:2.21.1'
testImplementation testLibs.spring.test
implementation libs.armeria.core
implementation libs.armeria.grpc
diff --git a/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java b/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java
index 18a711a8e6..e8c1c08563 100644
--- a/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java
+++ b/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java
@@ -79,7 +79,7 @@ void sending_alpha_and_beta_events_sends_to_sinks_that_take_both(final String so
inMemorySourceAccessor.submit(TESTING_KEY, allEvents);
- await().atMost(400, TimeUnit.MILLISECONDS)
+ await().atMost(2, TimeUnit.SECONDS)
.untilAsserted(() -> {
assertThat(inMemorySinkAccessor.get(sourceKeyToReceiveAll), not(empty()));
});
@@ -101,7 +101,7 @@ void sending_alpha_and_beta_events_sends_to_both_sinks() {
inMemorySourceAccessor.submit(TESTING_KEY, allEvents);
- await().atMost(400, TimeUnit.MILLISECONDS)
+ await().atMost(2, TimeUnit.SECONDS)
.untilAsserted(() -> {
assertThat(inMemorySinkAccessor.get(ALPHA_SOURCE_KEY), not(empty()));
assertThat(inMemorySinkAccessor.get(BETA_SOURCE_KEY), not(empty()));
@@ -136,9 +136,9 @@ void sending_non_alpha_beta_events_never_reaches_sink(final String sourceKey) th
inMemorySourceAccessor.submit(TESTING_KEY, randomEvents);
- Thread.sleep(1000);
-
- assertThat(inMemorySinkAccessor.get(sourceKey), empty());
+ await().during(1200, TimeUnit.MILLISECONDS)
+ .pollDelay(50, TimeUnit.MILLISECONDS)
+ .until(() -> inMemorySinkAccessor.get(sourceKey), empty());
}
private List> createEvents(final String value, final int numberToCreate) {
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java
index 21032873a4..2d32cb116c 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java
@@ -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;
@@ -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);
}
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java
index 1057418876..af9860cc9a 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java
@@ -6,7 +6,8 @@
package org.opensearch.dataprepper.acknowledgements;
import org.opensearch.dataprepper.model.event.EventHandle;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import java.util.concurrent.locks.ReentrantLock;
@@ -33,7 +34,12 @@ class AcknowledgementSetMonitor implements Runnable {
private final AtomicInteger numNullHandles;
private DefaultAcknowledgementSet getAcknowledgementSet(final EventHandle eventHandle) {
- return (DefaultAcknowledgementSet)((DefaultEventHandle)eventHandle).getAcknowledgementSet();
+ if (eventHandle instanceof DefaultEventHandle) {
+ InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+ return (DefaultAcknowledgementSet)internalEventHandle.getAcknowledgementSet();
+ } else {
+ throw new RuntimeException("Unsupported event handle");
+ }
}
public AcknowledgementSetMonitor() {
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java
index 3c8fe12159..1c3ef6032d 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java
@@ -5,10 +5,12 @@
package org.opensearch.dataprepper.acknowledgements;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -17,45 +19,73 @@
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 callback;
+ private Consumer 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 pendingAcknowledgments;
private Future callbackFuture;
private final DefaultAcknowledgementSetMetrics metrics;
+ private ScheduledFuture progressCheckFuture;
private boolean completed;
+ private AtomicInteger totalEventsAdded;
- public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer callback, final Duration expiryTime, final DefaultAcknowledgementSetMetrics metrics) {
+ public DefaultAcknowledgementSet(final ScheduledExecutorService scheduledExecutor,
+ final Consumer 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 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(new DefaultProgressCheck((double)numberOfEventsPending/totalEventsAdded.get()));
+ }
+ }
+
@Override
public void add(Event event) {
lock.lock();
try {
if (event instanceof JacksonEvent) {
- EventHandle eventHandle = new DefaultEventHandle(this);
- ((JacksonEvent) event).setEventHandle(eventHandle);
- pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
+ EventHandle eventHandle = event.getEventHandle();
+ if (eventHandle instanceof DefaultEventHandle) {
+ InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+ internalEventHandle.setAcknowledgementSet(this);
+ pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
+ totalEventsAdded.incrementAndGet();
+ }
}
} finally {
lock.unlock();
@@ -84,6 +114,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;
@@ -108,7 +141,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();
@@ -132,10 +168,13 @@ 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");
+ LOG.debug("Acknowledgement set is not completed. Delaying callback until it is completed");
}
}
} finally {
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java
index 104945960e..3f2e3761bd 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java
@@ -15,27 +15,27 @@
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");
@@ -43,7 +43,7 @@ public DefaultAcknowledgementSetManager(final ExecutorService callbackExecutor,
}
public AcknowledgementSet create(final Consumer 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;
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java
new file mode 100644
index 0000000000..87b7a8226d
--- /dev/null
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.acknowledgements;
+
+import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck;
+
+public class DefaultProgressCheck implements ProgressCheck {
+ double ratio;
+
+ public DefaultProgressCheck(double ratio) {
+ this.ratio = ratio;
+ }
+
+ @Override
+ public Double getRatio() {
+ return ratio;
+ }
+}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java
index 19fb683dcd..c8d6393682 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java
@@ -5,12 +5,14 @@
package org.opensearch.dataprepper.breaker;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.parser.model.CircuitBreakerConfig;
import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import java.util.List;
+import java.util.Optional;
/**
* The application config for circuit breakers. Used for wiring beans
@@ -34,4 +36,9 @@ InnerCircuitBreaker heapCircuitBreaker(final DataPrepperConfiguration dataPreppe
return null;
}
}
+
+ @Bean
+ public Optional circuitBreaker(final CircuitBreakerManager circuitBreakerManager) {
+ return circuitBreakerManager.getGlobalCircuitBreaker();
+ }
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java
index 23360ab04c..405b9e3bdc 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java
@@ -5,6 +5,8 @@
package org.opensearch.dataprepper.breaker;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
+
import java.util.List;
import java.util.Optional;
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java
index dc34bd5cd0..9993b9e4a8 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java
@@ -6,6 +6,7 @@
package org.opensearch.dataprepper.breaker;
import io.micrometer.core.instrument.Metrics;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.parser.model.HeapCircuitBreakerConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java
index 50fb2450cb..ebaa81fa7d 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java
@@ -5,6 +5,8 @@
package org.opensearch.dataprepper.breaker;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
+
/**
* Interface to signal that this {@link CircuitBreaker} to prevent
* access beyond the {@link CircuitBreakerManager}.
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java
deleted file mode 100644
index 025b130bf6..0000000000
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.dataprepper.event;
-
-import org.opensearch.dataprepper.model.event.EventHandle;
-import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
-import java.lang.ref.WeakReference;
-
-public class DefaultEventHandle implements EventHandle {
- private final WeakReference acknowledgementSetRef;
- public DefaultEventHandle(AcknowledgementSet acknowledgementSet) {
- this.acknowledgementSetRef = new WeakReference<>(acknowledgementSet);
- }
-
- public AcknowledgementSet getAcknowledgementSet() {
- return acknowledgementSetRef.get();
- }
-
- @Override
- public void release(boolean result) {
- AcknowledgementSet acknowledgementSet = getAcknowledgementSet();
- if (acknowledgementSet != null) {
- acknowledgementSet.release(this, result);
- }
- }
-}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java
index 0e63c682c4..d61f183f05 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java
@@ -5,14 +5,12 @@
package org.opensearch.dataprepper.parser;
-import org.opensearch.dataprepper.breaker.CircuitBreaker;
-import org.opensearch.dataprepper.model.CheckpointState;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.buffer.DelegatingBuffer;
import org.opensearch.dataprepper.model.record.Record;
-import java.time.Duration;
import java.util.Collection;
-import java.util.Map;
import java.util.concurrent.TimeoutException;
import static java.util.Objects.requireNonNull;
@@ -24,8 +22,7 @@
* @param The type of record.
* @since 2.1
*/
-class CircuitBreakingBuffer> implements Buffer {
- private final Buffer buffer;
+class CircuitBreakingBuffer> extends DelegatingBuffer implements Buffer {
private final CircuitBreaker circuitBreaker;
/**
@@ -35,7 +32,7 @@ class CircuitBreakingBuffer> implements Buffer {
* @param circuitBreaker The circuit breaker to check
*/
public CircuitBreakingBuffer(final Buffer buffer, final CircuitBreaker circuitBreaker) {
- this.buffer = requireNonNull(buffer);
+ super(buffer);
this.circuitBreaker = requireNonNull(circuitBreaker);
}
@@ -43,43 +40,25 @@ public CircuitBreakingBuffer(final Buffer buffer, final CircuitBreaker circui
public void write(final T record, final int timeoutInMillis) throws TimeoutException {
checkBreaker();
- buffer.write(record, timeoutInMillis);
+ super.write(record, timeoutInMillis);
}
@Override
public void writeAll(final Collection records, final int timeoutInMillis) throws Exception {
checkBreaker();
- buffer.writeAll(records, timeoutInMillis);
- }
-
- private void checkBreaker() throws TimeoutException {
- if(circuitBreaker.isOpen())
- throw new TimeoutException("Circuit breaker is open. Unable to write to buffer.");
- }
-
- @Override
- public Map.Entry, CheckpointState> read(final int timeoutInMillis) {
- return buffer.read(timeoutInMillis);
- }
-
- @Override
- public void checkpoint(final CheckpointState checkpointState) {
- buffer.checkpoint(checkpointState);
+ super.writeAll(records, timeoutInMillis);
}
@Override
- public boolean isEmpty() {
- return buffer.isEmpty();
- }
+ public void writeBytes(final byte[] bytes, final String key, final int timeoutInMillis) throws Exception {
+ checkBreaker();
- @Override
- public Duration getDrainTimeout() {
- return buffer.getDrainTimeout();
+ super.writeBytes(bytes, key, timeoutInMillis);
}
- @Override
- public void shutdown() {
- buffer.shutdown();
+ private void checkBreaker() throws TimeoutException {
+ if(circuitBreaker.isOpen())
+ throw new TimeoutException("Circuit breaker is open. Unable to write to buffer.");
}
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java
new file mode 100644
index 0000000000..9fb7a7a201
--- /dev/null
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.parser;
+
+import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.buffer.DelegatingBuffer;
+import org.opensearch.dataprepper.model.record.Record;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Buffer decorator created for pipelines that make use of multiple buffers, such as PeerForwarder-enabled pipelines. The decorator
+ * acts as a pass-through to the primary buffer for most methods except those that rely on a combination of the primary
+ * and second. For example, isEmpty depends on all buffers being empty.
+ *
+ * @since 2.0
+ */
+class MultiBufferDecorator> extends DelegatingBuffer implements Buffer {
+ private final List allBuffers;
+
+ MultiBufferDecorator(final Buffer primaryBuffer, final List secondaryBuffers) {
+ super(primaryBuffer);
+ allBuffers = new ArrayList<>(1 + secondaryBuffers.size());
+ allBuffers.add(primaryBuffer);
+ allBuffers.addAll(secondaryBuffers);
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return allBuffers.stream().allMatch(Buffer::isEmpty);
+ }
+
+ @Override
+ public Duration getDrainTimeout() {
+ return allBuffers.stream()
+ .map(Buffer::getDrainTimeout)
+ .reduce(Duration.ZERO, Duration::plus);
+ }
+
+ @Override
+ public void shutdown() {
+ allBuffers.forEach(Buffer::shutdown);
+ }
+}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java
index d520fe5062..64944aa47a 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java
@@ -28,7 +28,6 @@
import org.opensearch.dataprepper.pipeline.PipelineConnector;
import org.opensearch.dataprepper.pipeline.router.Router;
import org.opensearch.dataprepper.pipeline.router.RouterFactory;
-import org.opensearch.dataprepper.plugins.MultiBufferDecorator;
import org.opensearch.dataprepper.sourcecoordination.SourceCoordinatorFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -147,15 +146,7 @@ private void buildPipelineFromConfiguration(
final MultiBufferDecorator multiBufferDecorator = new MultiBufferDecorator(pipelineDefinedBuffer, secondaryBuffers);
- final Buffer buffer;
- if(source instanceof PipelineConnector) {
- buffer = multiBufferDecorator;
- } else {
- buffer = circuitBreakerManager.getGlobalCircuitBreaker()
- .map(circuitBreaker -> new CircuitBreakingBuffer<>(multiBufferDecorator, circuitBreaker))
- .map(b -> (Buffer)b)
- .orElseGet(() -> multiBufferDecorator);
- }
+ final Buffer buffer = applyCircuitBreakerToBuffer(source, multiBufferDecorator);
final Router router = routerFactory.createRouter(pipelineConfiguration.getRoutes());
@@ -313,4 +304,17 @@ List getSecondaryBuffers() {
.map(innerEntry -> innerEntry.getValue())
.collect(Collectors.toList());
}
+
+ private Buffer applyCircuitBreakerToBuffer(final Source source, final Buffer buffer) {
+ if (source instanceof PipelineConnector)
+ return buffer;
+
+ if(buffer.isWrittenOffHeapOnly())
+ return buffer;
+
+ return circuitBreakerManager.getGlobalCircuitBreaker()
+ .map(circuitBreaker -> new CircuitBreakingBuffer<>(buffer, circuitBreaker))
+ .map(b -> (Buffer) b)
+ .orElseGet(() -> buffer);
+ }
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java
index 5ddca5f10e..1b86adebfe 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java
@@ -23,6 +23,7 @@
import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesResponse;
import software.amazon.awssdk.services.servicediscovery.model.HttpInstanceSummary;
+import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Objects;
@@ -48,7 +49,7 @@ class AwsCloudMapPeerListProvider implements PeerListProvider, AutoCloseable {
private final String serviceName;
private final Map queryParameters;
private final AwsCloudMapDynamicEndpointGroup endpointGroup;
- private final int timeToRefreshSeconds;
+ private final Duration timeToRefresh;
private final Backoff backoff;
private final EventLoop eventLoop;
private final String domainName;
@@ -58,18 +59,18 @@ class AwsCloudMapPeerListProvider implements PeerListProvider, AutoCloseable {
final String namespaceName,
final String serviceName,
final Map queryParameters,
- final int timeToRefreshSeconds,
+ final Duration timeToRefresh,
final Backoff backoff,
final PluginMetrics pluginMetrics) {
this.awsServiceDiscovery = Objects.requireNonNull(awsServiceDiscovery);
this.namespaceName = Objects.requireNonNull(namespaceName);
this.serviceName = Objects.requireNonNull(serviceName);
this.queryParameters = Objects.requireNonNull(queryParameters);
- this.timeToRefreshSeconds = timeToRefreshSeconds;
+ this.timeToRefresh = timeToRefresh;
this.backoff = Objects.requireNonNull(backoff);
- if (timeToRefreshSeconds < 1)
- throw new IllegalArgumentException("timeToRefreshSeconds must be positive. Actual: " + timeToRefreshSeconds);
+ if (timeToRefresh.isNegative() || timeToRefresh.isZero())
+ throw new IllegalArgumentException("timeToRefreshSeconds must be positive. Actual: " + timeToRefresh);
eventLoop = CommonPools.workerGroup().next();
LOG.info("Using AWS CloudMap for Peer Forwarding. namespace='{}', serviceName='{}'",
@@ -92,7 +93,7 @@ static AwsCloudMapPeerListProvider createPeerListProvider(final PeerForwarderCon
final Map queryParameters = peerForwarderConfiguration.getAwsCloudMapQueryParameters();
final Backoff standardBackoff = Backoff.exponential(ONE_SECOND, TWENTY_SECONDS).withJitter(TWENTY_PERCENT);
- final int timeToRefreshSeconds = 20;
+ final Duration timeToRefresh = Duration.ofSeconds(20);
final PluginMetrics awsSdkMetrics = PluginMetrics.fromNames("sdk", "aws");
@@ -108,7 +109,7 @@ static AwsCloudMapPeerListProvider createPeerListProvider(final PeerForwarderCon
namespace,
serviceName,
queryParameters,
- timeToRefreshSeconds,
+ timeToRefresh,
standardBackoff,
pluginMetrics);
}
@@ -175,7 +176,7 @@ private void discoverInstances() {
LOG.warn("Failed to update endpoints.", ex);
} finally {
scheduledDiscovery = eventLoop.schedule(this::discoverInstances,
- timeToRefreshSeconds, TimeUnit.SECONDS);
+ timeToRefresh.toMillis(), TimeUnit.MILLISECONDS);
}
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java
index 0bd5d469ca..a71f8d14b1 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java
@@ -344,8 +344,11 @@ List> publishToSinks(final Collection records) {
InactiveAcknowledgementSetManager.getInstance(),
sinks);
router.route(records, sinks, getRecordStrategy, (sink, events) ->
- sinkFutures.add(sinkExecutorService.submit(() -> sink.output(events), null))
- );
+ sinkFutures.add(sinkExecutorService.submit(() -> {
+ sink.updateLatencyMetrics(events);
+ sink.output(events);
+ }, null))
+ );
return sinkFutures;
}
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java
index 56d81ba68b..fb4effb413 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java
@@ -13,6 +13,8 @@
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
import org.opensearch.dataprepper.pipeline.common.FutureHelper;
import org.opensearch.dataprepper.pipeline.common.FutureHelperResult;
import org.slf4j.Logger;
@@ -22,7 +24,6 @@
import java.util.List;
import java.util.ArrayList;
import java.util.Map;
-import java.util.Objects;
import java.util.Set;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
@@ -97,10 +98,15 @@ private void processAcknowledgements(List inputEvents, Collection outputR
// For each event in the input events list that is not present in the output events, send positive acknowledgement, if acknowledgements are enabled for it
inputEvents.forEach(event -> {
EventHandle eventHandle = event.getEventHandle();
- if (Objects.nonNull(eventHandle) && !outputEventsSet.contains(event)) {
- eventHandle.release(true);
- } else if (acknowledgementsEnabled && Objects.isNull(eventHandle)) {
- invalidEventHandlesCounter.increment();
+ if (eventHandle != null && eventHandle instanceof DefaultEventHandle) {
+ InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+ if (internalEventHandle.getAcknowledgementSet() != null && !outputEventsSet.contains(event)) {
+ eventHandle.release(true);
+ } else if (acknowledgementsEnabled) {
+ invalidEventHandlesCounter.increment();
+ }
+ } else if (eventHandle != null) {
+ throw new RuntimeException("Unexpected EventHandle");
}
});
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java
index df7316b981..1bd2944c2e 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java
@@ -18,7 +18,7 @@
import org.opensearch.dataprepper.model.event.EventHandle;
import org.opensearch.dataprepper.model.event.EventBuilder;
import org.opensearch.dataprepper.model.event.EventMetadata;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
import org.opensearch.dataprepper.acknowledgements.InactiveAcknowledgementSetManager;
@@ -65,7 +65,7 @@ private void acquireEventReference(final Record record) {
}
if (referencedRecords.contains(record) || ((routedRecords != null) && routedRecords.contains(record))) {
EventHandle eventHandle = ((JacksonEvent)record.getData()).getEventHandle();
- if (eventHandle != null) {
+ if (eventHandle != null && eventHandle instanceof DefaultEventHandle) {
acknowledgementSetManager.acquireEventReference(eventHandle);
}
} else if (!referencedRecords.contains(record)) {
@@ -97,7 +97,7 @@ public Record getRecord(final Record record) {
JacksonEvent newRecordEvent;
Record newRecord;
DefaultEventHandle eventHandle = (DefaultEventHandle)recordEvent.getEventHandle();
- if (eventHandle != null) {
+ if (eventHandle != null && eventHandle.getAcknowledgementSet() != null) {
final EventMetadata eventMetadata = recordEvent.getMetadata();
final EventBuilder eventBuilder = (EventBuilder) eventFactory.eventBuilder(EventBuilder.class).withEventMetadata(eventMetadata).withData(recordEvent.toMap());
newRecordEvent = (JacksonEvent) eventBuilder.build();
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java
index bb688abb8f..6be9244c08 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java
@@ -5,6 +5,7 @@
package org.opensearch.dataprepper.plugin;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
import org.opensearch.dataprepper.model.sink.SinkContext;
import org.opensearch.dataprepper.metrics.PluginMetrics;
@@ -74,6 +75,8 @@ private ComponentPluginArgumentsContext(final Builder builder) {
if (builder.sinkContext != null) {
typedArgumentsSuppliers.put(SinkContext.class, () -> builder.sinkContext);
}
+
+ typedArgumentsSuppliers.put(CircuitBreaker.class, () -> builder.circuitBreaker);
}
@Override
@@ -135,6 +138,7 @@ static class Builder {
private AcknowledgementSetManager acknowledgementSetManager;
private PluginConfigObservable pluginConfigObservable;
private SinkContext sinkContext;
+ private CircuitBreaker circuitBreaker;
Builder withPluginConfiguration(final Object pluginConfiguration) {
this.pluginConfiguration = pluginConfiguration;
@@ -181,6 +185,11 @@ Builder withPluginConfigurationObservable(final PluginConfigObservable pluginCon
return this;
}
+ Builder withCircuitBreaker(final CircuitBreaker circuitBreaker) {
+ this.circuitBreaker = circuitBreaker;
+ return this;
+ }
+
ComponentPluginArgumentsContext build() {
return new ComponentPluginArgumentsContext(this);
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java
index bd53b9678a..8e6f070526 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java
@@ -5,6 +5,7 @@
package org.opensearch.dataprepper.plugin;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.plugin.PluginConfigObservable;
import org.opensearch.dataprepper.model.sink.SinkContext;
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
@@ -15,6 +16,7 @@
import org.opensearch.dataprepper.acknowledgements.DefaultAcknowledgementSetManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.DependsOn;
import javax.inject.Inject;
@@ -43,6 +45,7 @@ public class DefaultPluginFactory implements PluginFactory {
private final DefaultEventFactory eventFactory;
private final DefaultAcknowledgementSetManager acknowledgementSetManager;
private final PluginConfigurationObservableFactory pluginConfigurationObservableFactory;
+ private final CircuitBreaker circuitBreaker;
@Inject
DefaultPluginFactory(
@@ -52,8 +55,10 @@ public class DefaultPluginFactory implements PluginFactory {
final PluginBeanFactoryProvider pluginBeanFactoryProvider,
final DefaultEventFactory eventFactory,
final DefaultAcknowledgementSetManager acknowledgementSetManager,
- final PluginConfigurationObservableFactory pluginConfigurationObservableFactory
- ) {
+ final PluginConfigurationObservableFactory pluginConfigurationObservableFactory,
+ @Autowired(required = false) final CircuitBreaker circuitBreaker
+ ) {
+ this.circuitBreaker = circuitBreaker;
Objects.requireNonNull(pluginProviderLoader);
Objects.requireNonNull(pluginConfigurationObservableFactory);
this.pluginCreator = Objects.requireNonNull(pluginCreator);
@@ -131,6 +136,7 @@ private ComponentPluginArgumentsContext getConstructionContext(final PluginS
.withAcknowledgementSetManager(acknowledgementSetManager)
.withPluginConfigurationObservable(pluginConfigObservable)
.withSinkContext(sinkContext)
+ .withCircuitBreaker(circuitBreaker)
.build();
}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java
deleted file mode 100644
index cd199a7486..0000000000
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java
+++ /dev/null
@@ -1,83 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.dataprepper.plugins;
-
-import org.opensearch.dataprepper.model.CheckpointState;
-import org.opensearch.dataprepper.model.buffer.Buffer;
-import org.opensearch.dataprepper.model.record.Record;
-
-import java.time.Duration;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.TimeoutException;
-import java.util.stream.Stream;
-
-/**
- * Buffer decorator created for pipelines that make use of multiple buffers, such as PeerForwarder-enabled pipelines. The decorator
- * acts as a pass-through to the primary buffer for all methods except isEmpty, which verifies that all buffers are empty.
- *
- * @since 2.0
- */
-public class MultiBufferDecorator> implements Buffer {
- private final Buffer primaryBuffer;
- private final List secondaryBuffers;
-
- public MultiBufferDecorator(final Buffer primaryBuffer, final List secondaryBuffers) {
- this.primaryBuffer = primaryBuffer;
- this.secondaryBuffers = secondaryBuffers;
- }
-
- @Override
- public void write(final T record, final int timeoutInMillis) throws TimeoutException {
- primaryBuffer.write(record, timeoutInMillis);
- }
-
- @Override
- public void writeAll(final Collection records, final int timeoutInMillis) throws Exception {
- primaryBuffer.writeAll(records, timeoutInMillis);
- }
-
- @Override
- public void writeBytes(final byte[] bytes, final String key, int timeoutInMillis) throws Exception {
- primaryBuffer.writeBytes(bytes, key, timeoutInMillis);
- }
-
- @Override
- public Map.Entry, CheckpointState> read(final int timeoutInMillis) {
- return primaryBuffer.read(timeoutInMillis);
- }
-
- @Override
- public void checkpoint(final CheckpointState checkpointState) {
- primaryBuffer.checkpoint(checkpointState);
- }
-
- @Override
- public boolean isByteBuffer() {
- return primaryBuffer.isByteBuffer();
- }
-
- @Override
- public boolean isEmpty() {
- return primaryBuffer.isEmpty() && secondaryBuffers.stream()
- .map(Buffer::isEmpty)
- .allMatch(result -> result == true);
- }
-
- @Override
- public Duration getDrainTimeout() {
- return Stream.concat(Stream.of(primaryBuffer), secondaryBuffers.stream())
- .map(Buffer::getDrainTimeout)
- .reduce(Duration.ZERO, Duration::plus);
- }
-
- @Override
- public void shutdown() {
- primaryBuffer.shutdown();
- secondaryBuffers.forEach(Buffer::shutdown);
- }
-}
diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java
index 7ada401383..5d4fb3ae77 100644
--- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java
+++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java
@@ -127,7 +127,7 @@ public Optional acquireAvailablePartition(String partit
@Override
- public void saveProgressStateForPartition(EnhancedSourcePartition partition) {
+ public void saveProgressStateForPartition(EnhancedSourcePartition partition, final Duration ownershipTimeoutRenewal) {
String partitionType = partition.getPartitionType() == null ? DEFAULT_GLOBAL_STATE_PARTITION_TYPE : partition.getPartitionType();
LOG.debug("Try to save progress for partition {} (Type {})", partition.getPartitionKey(), partitionType);
@@ -140,7 +140,7 @@ public void saveProgressStateForPartition(EnhancedSourcePartition partiti
final SourcePartitionStoreItem updateItem = partition.getSourcePartitionStoreItem();
// Also extend the timeout of the lease (ownership)
if (updateItem.getPartitionOwnershipTimeout() != null) {
- updateItem.setPartitionOwnershipTimeout(Instant.now().plus(DEFAULT_LEASE_TIMEOUT));
+ updateItem.setPartitionOwnershipTimeout(Instant.now().plus(ownershipTimeoutRenewal == null ? DEFAULT_LEASE_TIMEOUT : ownershipTimeoutRenewal));
}
updateItem.setPartitionProgressState(partition.convertPartitionProgressStatetoString(partition.getProgressState()));
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java
index 958a3331b3..c3051c7fc0 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java
@@ -5,14 +5,14 @@
package org.opensearch.dataprepper;
-import org.opensearch.dataprepper.model.configuration.PluginModel;
-import org.opensearch.dataprepper.model.configuration.PluginSetting;
-import org.opensearch.dataprepper.model.configuration.SinkModel;
-import org.opensearch.dataprepper.parser.model.PipelineConfiguration;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+import org.opensearch.dataprepper.model.configuration.PluginModel;
+import org.opensearch.dataprepper.model.configuration.PluginSetting;
+import org.opensearch.dataprepper.model.configuration.SinkModel;
+import org.opensearch.dataprepper.parser.model.PipelineConfiguration;
import java.io.File;
import java.io.IOException;
@@ -38,6 +38,7 @@ public class TestDataProvider {
public static final String VALID_MULTIPLE_PIPELINE_CONFIG_FILE = "src/test/resources/valid_multiple_pipeline_configuration.yml";
public static final String VALID_PIPELINE_CONFIG_FILE_WITH_EXTENSIONS = "src/test/resources/valid_pipeline_configuration_with_extensions.yml";
public static final String VALID_SINGLE_PIPELINE_EMPTY_SOURCE_PLUGIN_FILE = "src/test/resources/single_pipeline_valid_empty_source_plugin_settings.yml";
+ public static final String VALID_OFF_HEAP_FILE = "src/test/resources/single_pipeline_valid_off_heap_buffer.yml";
public static final String CONNECTED_PIPELINE_ROOT_SOURCE_INCORRECT = "src/test/resources/connected_pipeline_incorrect_root_source.yml";
public static final String CONNECTED_PIPELINE_CHILD_PIPELINE_INCORRECT = "src/test/resources/connected_pipeline_incorrect_child_pipeline.yml";
public static final String CYCLE_MULTIPLE_PIPELINE_CONFIG_FILE = "src/test/resources/cyclic_multiple_pipeline_configuration.yml";
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java
index 8c9d065704..6c85b5c4de 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java
@@ -5,7 +5,7 @@
package org.opensearch.dataprepper.acknowledgements;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mock;
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java
index 486617e9a0..1b87d6c849 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java
@@ -6,7 +6,8 @@
package org.opensearch.dataprepper.acknowledgements;
import org.opensearch.dataprepper.model.event.JacksonEvent;
-import org.opensearch.dataprepper.model.event.EventHandle;
+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;
@@ -18,20 +19,18 @@
import static org.awaitility.Awaitility.await;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.ArgumentMatchers.any;
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;
@@ -40,32 +39,35 @@ class DefaultAcknowledgementSetManagerTests {
@Mock
JacksonEvent event3;
- EventHandle eventHandle1;
- EventHandle eventHandle2;
- EventHandle eventHandle3;
- Boolean result;
+ private PluginMetrics pluginMetrics;
+ private DefaultEventHandle eventHandle1;
+ private DefaultEventHandle eventHandle2;
+ private DefaultEventHandle eventHandle3;
+ private DefaultEventHandle eventHandle4;
+ private DefaultEventHandle eventHandle5;
+ private DefaultEventHandle eventHandle6;
+ private Boolean result;
+ private double currentRatio;
@BeforeEach
void setup() {
- callbackExecutor = Executors.newFixedThreadPool(2);
+ currentRatio = 0;
+ callbackExecutor = Executors.newScheduledThreadPool(2);
event1 = mock(JacksonEvent.class);
- doAnswer((i) -> {
- eventHandle1 = i.getArgument(0);
- return null;
- }).when(event1).setEventHandle(any());
+ eventHandle1 = mock(DefaultEventHandle.class);
lenient().when(event1.getEventHandle()).thenReturn(eventHandle1);
+ pluginMetrics = mock(PluginMetrics.class);
event2 = mock(JacksonEvent.class);
- doAnswer((i) -> {
- eventHandle2 = i.getArgument(0);
- return null;
- }).when(event2).setEventHandle(any());
+ eventHandle2 = mock(DefaultEventHandle.class);
lenient().when(event2.getEventHandle()).thenReturn(eventHandle2);
acknowledgementSetManager = createObjectUnderTest();
AcknowledgementSet acknowledgementSet1 = acknowledgementSetManager.create((flag) -> { result = flag; }, TEST_TIMEOUT);
acknowledgementSet1.add(event1);
acknowledgementSet1.add(event2);
+ lenient().when(eventHandle1.getAcknowledgementSet()).thenReturn(acknowledgementSet1);
+ lenient().when(eventHandle2.getAcknowledgementSet()).thenReturn(acknowledgementSet1);
acknowledgementSet1.complete();
}
@@ -82,8 +84,6 @@ void testBasic() {
assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
assertThat(result, equalTo(true));
});
- assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0));
- assertThat(result, equalTo(true));
}
@Test
@@ -91,20 +91,21 @@ 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
void testMultipleAcknowledgementSets() {
event3 = mock(JacksonEvent.class);
- doAnswer((i) -> {
- eventHandle3 = i.getArgument(0);
- return null;
- }).when(event3).setEventHandle(any());
+ eventHandle3 = mock(DefaultEventHandle.class);
lenient().when(event3.getEventHandle()).thenReturn(eventHandle3);
AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, TEST_TIMEOUT);
acknowledgementSet2.add(event3);
+ lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2);
acknowledgementSet2.complete();
acknowledgementSetManager.releaseEventReference(eventHandle2, true);
@@ -114,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((progressCheck) -> {currentRatio = progressCheck.getRatio();}, 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((progressCheck) -> {currentRatio = progressCheck.getRatio();}, 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));
+ });
+
+ }
+
}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java
index c4403e4b2f..5deba46be8 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java
@@ -6,16 +6,18 @@
package org.opensearch.dataprepper.acknowledgements;
import org.awaitility.Awaitility;
+import static org.awaitility.Awaitility.await;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.event.JacksonEvent;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import java.time.Duration;
-import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
@@ -24,7 +26,6 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
@@ -34,9 +35,16 @@ class DefaultAcknowledgementSetTests {
private DefaultAcknowledgementSet defaultAcknowledgementSet;
@Mock
private JacksonEvent event;
+ @Mock
+ private JacksonEvent event2;
+ @Mock
private DefaultEventHandle handle;
+ @Mock
+ private DefaultEventHandle handle2;
+
+ private double currentRatio;
- private ExecutorService executor;
+ private ScheduledExecutorService executor;
private Boolean acknowledgementSetResult;
private final Duration TEST_TIMEOUT = Duration.ofMillis(5000);
private AtomicBoolean callbackInterrupted;
@@ -71,19 +79,28 @@ private DefaultAcknowledgementSet createObjectUnderTestWithCallback(Consumer {
- handle = i.getArgument(0);
+ handle = mock(DefaultEventHandle.class);
+ lenient().doAnswer(a -> {
+ AcknowledgementSet acknowledgementSet = a.getArgument(0);
+ lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
return null;
- }).when(event).setEventHandle(any());
+ }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class));
lenient().when(event.getEventHandle()).thenReturn(handle);
+ event2 = mock(JacksonEvent.class);
+ lenient().doAnswer(a -> {
+ AcknowledgementSet acknowledgementSet = a.getArgument(0);
+ lenient().when(handle2.getAcknowledgementSet()).thenReturn(acknowledgementSet);
+ return null;
+ }).when(handle2).setAcknowledgementSet(any(AcknowledgementSet.class));
+ handle2 = mock(DefaultEventHandle.class);
+ lenient().when(event2.getEventHandle()).thenReturn(handle2);
}
@Test
@@ -115,7 +132,6 @@ void testDefaultAcknowledgementInvalidAcquire() {
defaultAcknowledgementSet.add(event);
defaultAcknowledgementSet.complete();
DefaultAcknowledgementSet secondAcknowledgementSet = createObjectUnderTest();
- DefaultEventHandle handle2 = new DefaultEventHandle(secondAcknowledgementSet);
defaultAcknowledgementSet.acquire(handle2);
assertThat(invalidAcquiresCounter, equalTo(1));
}
@@ -125,7 +141,6 @@ void testDefaultAcknowledgementInvalidRelease() {
defaultAcknowledgementSet.add(event);
defaultAcknowledgementSet.complete();
DefaultAcknowledgementSet secondAcknowledgementSet = createObjectUnderTest();
- DefaultEventHandle handle2 = new DefaultEventHandle(secondAcknowledgementSet);
assertThat(defaultAcknowledgementSet.release(handle2, true), equalTo(false));
assertThat(invalidReleasesCounter, equalTo(1));
}
@@ -170,6 +185,11 @@ void testDefaultAcknowledgementSetNegativeAcknowledgements() throws Exception {
defaultAcknowledgementSet.add(event);
defaultAcknowledgementSet.complete();
assertThat(handle, not(equalTo(null)));
+ lenient().doAnswer(a -> {
+ AcknowledgementSet acknowledgementSet = a.getArgument(0);
+ lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
+ return null;
+ }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class));
assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet));
defaultAcknowledgementSet.acquire(handle);
assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(false));
@@ -198,6 +218,11 @@ void testDefaultAcknowledgementSetExpirations() throws Exception {
);
defaultAcknowledgementSet.add(event);
defaultAcknowledgementSet.complete();
+ lenient().doAnswer(a -> {
+ AcknowledgementSet acknowledgementSet = a.getArgument(0);
+ lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
+ return null;
+ }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class));
assertThat(handle, not(equalTo(null)));
assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet));
assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(true));
@@ -210,4 +235,46 @@ void testDefaultAcknowledgementSetExpirations() throws Exception {
.until(() -> callbackInterrupted.get());
assertThat(callbackInterrupted.get(), equalTo(true));
}
+
+ @Test
+ void testDefaultAcknowledgementSetWithProgressCheck() throws Exception {
+ defaultAcknowledgementSet = createObjectUnderTestWithCallback(
+ (flag) -> {
+ acknowledgementSetResult = flag;
+ }
+ );
+ defaultAcknowledgementSet.addProgressCheck(
+ (progressCheck) -> {
+ currentRatio = progressCheck.getRatio();
+ },
+ Duration.ofSeconds(1)
+ );
+ defaultAcknowledgementSet.add(event);
+ defaultAcknowledgementSet.add(event2);
+ defaultAcknowledgementSet.complete();
+ lenient().doAnswer(a -> {
+ AcknowledgementSet acknowledgementSet = a.getArgument(0);
+ lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
+ return null;
+ }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class));
+ assertThat(handle, not(equalTo(null)));
+ assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet));
+ await().atMost(Duration.ofSeconds(5))
+ .untilAsserted(() -> {
+ assertThat(currentRatio, equalTo(1.0));
+ });
+ assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(false));
+ await().atMost(Duration.ofSeconds(5))
+ .untilAsserted(() -> {
+ assertThat(currentRatio, equalTo(0.5));
+ });
+ assertThat(defaultAcknowledgementSet.release(handle2, true), equalTo(true));
+ Awaitility.waitAtMost(Duration.ofSeconds(10))
+ .pollDelay(Duration.ofMillis(500))
+ .until(() -> defaultAcknowledgementSet.isDone());
+ await().atMost(Duration.ofSeconds(5))
+ .untilAsserted(() -> {
+ assertThat(acknowledgementSetResult, equalTo(true));
+ });
+ }
}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java
index ff14013db3..44c32b4f4d 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java
@@ -13,6 +13,7 @@
import org.junit.jupiter.params.provider.CsvSource;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.types.ByteCount;
import org.opensearch.dataprepper.parser.model.CircuitBreakerConfig;
import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration;
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java
index a57a83caa1..e92bf39e61 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java
@@ -13,6 +13,7 @@
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import java.util.Collections;
import java.util.List;
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java
deleted file mode 100644
index 2b0895ad6c..0000000000
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java
+++ /dev/null
@@ -1,32 +0,0 @@
-/*
- * Copyright OpenSearch Contributors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package org.opensearch.dataprepper.event;
-
-import org.opensearch.dataprepper.model.event.EventHandle;
-import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.equalTo;
-import org.junit.jupiter.api.Test;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-import static org.mockito.Mockito.verify;
-import static org.mockito.ArgumentMatchers.any;
-import org.mockito.Mock;
-
-class DefaultEventHandleTests {
- @Mock
- private AcknowledgementSet acknowledgementSet;
-
- @Test
- void testBasic() {
- acknowledgementSet = mock(AcknowledgementSet.class);
- when(acknowledgementSet.release(any(EventHandle.class), any(Boolean.class))).thenReturn(true);
- DefaultEventHandle eventHandle = new DefaultEventHandle(acknowledgementSet);
- assertThat(eventHandle.getAcknowledgementSet(), equalTo(acknowledgementSet));
- eventHandle.release(true);
- verify(acknowledgementSet).release(eventHandle, true);
- }
-}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java
index 33e8832662..faf8c25120 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java
@@ -14,7 +14,7 @@
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
-import org.opensearch.dataprepper.breaker.CircuitBreaker;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.CheckpointState;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.record.Record;
@@ -25,6 +25,7 @@
import java.util.List;
import java.util.Map;
import java.util.Random;
+import java.util.UUID;
import java.util.concurrent.TimeoutException;
import static org.hamcrest.MatcherAssert.assertThat;
@@ -81,6 +82,14 @@ void shutdown_calls_buffer_shutdown() {
verify(buffer).shutdown();
}
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ void isByteBuffer_returns_value_of_inner_buffer(boolean innerIsByteBuffer) {
+ when(buffer.isByteBuffer()).thenReturn(innerIsByteBuffer);
+
+ assertThat(createObjectUnderTest().isByteBuffer(), equalTo(innerIsByteBuffer));
+ }
+
@Nested
class NoCircuitBreakerChecks {
@AfterEach
@@ -159,4 +168,38 @@ void writeAll_should_check_CircuitBreaker_and_throw_if_open() {
verify(circuitBreaker).isOpen();
}
}
+
+ @Nested
+ class WithBytes {
+ private byte[] bytes;
+
+ private String key;
+
+ @BeforeEach
+ void setUp() {
+ bytes = UUID.randomUUID().toString().getBytes();
+ key = UUID.randomUUID().toString();
+ }
+
+ @Test
+ void writeBytes_should_check_CircuitBreaker_and_call_inner_write_if_not_open() throws Exception {
+ when(circuitBreaker.isOpen()).thenReturn(false);
+
+ createObjectUnderTest().writeBytes(bytes, key, timeoutMillis);
+
+ verify(buffer).writeBytes(bytes, key, timeoutMillis);
+ verify(circuitBreaker).isOpen();
+ }
+
+ @Test
+ void writeBytes_should_check_CircuitBreaker_and_throw_if_open() {
+ when(circuitBreaker.isOpen()).thenReturn(true);
+
+ CircuitBreakingBuffer> objectUnderTest = createObjectUnderTest();
+ assertThrows(TimeoutException.class, () -> objectUnderTest.writeBytes(bytes, key, timeoutMillis));
+
+ verifyNoInteractions(buffer);
+ verify(circuitBreaker).isOpen();
+ }
+ }
}
\ No newline at end of file
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java
similarity index 99%
rename from data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java
rename to data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java
index 003e47c024..fa405154a1 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java
@@ -3,7 +3,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-package org.opensearch.dataprepper.plugins;
+package org.opensearch.dataprepper.parser;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java
index ce44b298fd..611dfaaf77 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java
@@ -16,7 +16,7 @@
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.opensearch.dataprepper.TestDataProvider;
-import org.opensearch.dataprepper.breaker.CircuitBreaker;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.breaker.CircuitBreakerManager;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.configuration.PipelinesDataFlowModel;
@@ -54,6 +54,7 @@
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
@@ -303,6 +304,27 @@ void parseConfiguration_uses_CircuitBreaking_buffer_when_circuit_breakers_applie
verify(dataPrepperConfiguration).getPipelineExtensions();
}
+ @Test
+ void parseConfiguration_uses_unwrapped_buffer_when_circuit_breakers_applied_but_Buffer_is_off_heap() {
+ final PipelineTransformer objectUnderTest =
+ createObjectUnderTest(TestDataProvider.VALID_OFF_HEAP_FILE);
+
+ final Map pipelineMap = objectUnderTest.transformConfiguration();
+
+ assertThat(pipelineMap.size(), equalTo(1));
+ assertThat(pipelineMap, hasKey("test-pipeline-1"));
+ final Pipeline pipeline = pipelineMap.get("test-pipeline-1");
+ assertThat(pipeline, notNullValue());
+ assertThat(pipeline.getBuffer(), notNullValue());
+ assertThat(pipeline.getBuffer(), CoreMatchers.not(instanceOf(CircuitBreakingBuffer.class)));
+
+ verifyNoInteractions(circuitBreakerManager);
+ verify(dataPrepperConfiguration).getProcessorShutdownTimeout();
+ verify(dataPrepperConfiguration).getSinkShutdownTimeout();
+ verify(dataPrepperConfiguration).getPeerForwarderConfiguration();
+ verify(dataPrepperConfiguration).getPipelineExtensions();
+ }
+
@Test
void parseConfiguration_uses_unwrapped_buffer_when_no_circuit_breakers_are_applied() {
when(circuitBreakerManager.getGlobalCircuitBreaker())
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java
index 0c7c23093a..67429e27ef 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java
@@ -5,11 +5,9 @@
package org.opensearch.dataprepper.peerforwarder.discovery;
-import org.opensearch.dataprepper.metrics.PluginMetrics;
import com.linecorp.armeria.client.Endpoint;
import com.linecorp.armeria.client.retry.Backoff;
import org.apache.commons.lang3.RandomStringUtils;
-import org.awaitility.Awaitility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Nested;
@@ -19,11 +17,13 @@
import org.mockito.ArgumentCaptor;
import org.mockito.ArgumentMatchers;
import org.mockito.InOrder;
+import org.opensearch.dataprepper.metrics.PluginMetrics;
import software.amazon.awssdk.services.servicediscovery.ServiceDiscoveryAsyncClient;
import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesRequest;
import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesResponse;
import software.amazon.awssdk.services.servicediscovery.model.HttpInstanceSummary;
+import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -38,6 +38,7 @@
import java.util.stream.Collectors;
import java.util.stream.IntStream;
+import static org.awaitility.Awaitility.await;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.notNullValue;
import static org.hamcrest.CoreMatchers.nullValue;
@@ -59,7 +60,7 @@ class AwsCloudMapPeerListProviderTest {
private String namespaceName;
private String serviceName;
private Map queryParameters;
- private int timeToRefreshSeconds;
+ private Duration timeToRefresh;
private Backoff backoff;
private PluginMetrics pluginMetrics;
private List objectsToClose;
@@ -71,7 +72,7 @@ void setUp() {
serviceName = RandomStringUtils.randomAlphabetic(10);
queryParameters = generateRandomStringMap();
- timeToRefreshSeconds = 1;
+ timeToRefresh = Duration.ofMillis(200);
backoff = mock(Backoff.class);
pluginMetrics = mock(PluginMetrics.class);
@@ -85,7 +86,7 @@ void tearDown() {
private AwsCloudMapPeerListProvider createObjectUnderTest() {
final AwsCloudMapPeerListProvider objectUnderTest =
- new AwsCloudMapPeerListProvider(awsServiceDiscovery, namespaceName, serviceName, queryParameters, timeToRefreshSeconds, backoff, pluginMetrics);
+ new AwsCloudMapPeerListProvider(awsServiceDiscovery, namespaceName, serviceName, queryParameters, timeToRefresh, backoff, pluginMetrics);
objectsToClose.add(objectUnderTest);
return objectUnderTest;
}
@@ -133,7 +134,7 @@ void constructor_throws_with_null_Backoff() {
@ParameterizedTest
@ValueSource(ints = {Integer.MIN_VALUE, -10, -1, 0})
void constructor_throws_with_non_positive_timeToRefreshSeconds(final int badTimeToRefresh) {
- timeToRefreshSeconds = badTimeToRefresh;
+ timeToRefresh = Duration.ofSeconds(badTimeToRefresh);
assertThrows(IllegalArgumentException.class,
this::createObjectUnderTest);
@@ -374,7 +375,7 @@ private void waitUntilDiscoverInstancesCalledAtLeastOnce() {
*/
private void waitUntilDiscoverInstancesCalledAtLeast(final int timesCalled) {
final long waitTimeMillis = (long) timesCalled * WAIT_TIME_MULTIPLIER_MILLIS;
- Awaitility.waitAtMost(waitTimeMillis, TimeUnit.MILLISECONDS)
+ await().atMost(waitTimeMillis, TimeUnit.MILLISECONDS)
.pollDelay(100, TimeUnit.MILLISECONDS)
.untilAsserted(() -> then(awsServiceDiscovery)
.should(atLeast(timesCalled))
@@ -388,7 +389,7 @@ private void waitUntilDiscoverInstancesCalledAtLeast(final int timesCalled) {
* @param objectUnderTest The object to wait for.
*/
private void waitUntilPeerListPopulated(final AwsCloudMapPeerListProvider objectUnderTest) {
- Awaitility.waitAtMost(2, TimeUnit.SECONDS)
+ await().atMost(5, TimeUnit.SECONDS)
.pollDelay(100, TimeUnit.MILLISECONDS)
.untilAsserted(() -> {
final List actualPeers = objectUnderTest.getPeerList();
@@ -411,7 +412,7 @@ private static Map generateRandomStringMap() {
final Map map = new HashMap<>();
IntStream.range(0, random.nextInt(5) + 1)
- .mapToObj(num -> map.put(UUID.randomUUID().toString(), UUID.randomUUID().toString()));
+ .forEach(num -> map.put(UUID.randomUUID().toString(), UUID.randomUUID().toString()));
return map;
}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java
index 2fc86b80eb..75c1154baa 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java
@@ -10,10 +10,11 @@
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.configuration.PluginSetting;
import org.opensearch.dataprepper.model.event.EventFactory;
-import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.processor.Processor;
import org.opensearch.dataprepper.model.record.Record;
@@ -83,7 +84,7 @@ class PipelineTests {
private Duration peerForwarderDrainTimeout;
private EventFactory eventFactory;
private JacksonEvent event;
- private EventHandle eventHandle;
+ private DefaultEventHandle eventHandle;
private AcknowledgementSetManager acknowledgementSetManager;
@BeforeEach
@@ -391,6 +392,7 @@ class PublishToSink {
private List records;
private List> dataFlowComponents;
private Source mockSource;
+ private AcknowledgementSet acknowledgementSet;
@BeforeEach
void setUp() {
@@ -426,8 +428,10 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM
RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2);
Record rec = records.get(0);
event = mock(JacksonEvent.class);
- eventHandle = mock(EventHandle.class);
+ eventHandle = mock(DefaultEventHandle.class);
+ acknowledgementSet = mock(AcknowledgementSet.class);
when(event.getEventHandle()).thenReturn(eventHandle);
+ when(eventHandle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
when(rec.getData()).thenReturn(event);
routerCopyRecordStrategy.getRecord(rec);
routerCopyRecordStrategy.getRecord(rec);
@@ -437,7 +441,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM
Pipeline pipeline = createObjectUnderTest();
when(mockSource.areAcknowledgementsEnabled()).thenReturn(true);
pipeline.publishToSinks(records);
- verify(acknowledgementSetManager).acquireEventReference(any(EventHandle.class));
+ verify(acknowledgementSetManager).acquireEventReference(any(DefaultEventHandle.class));
verify(router)
.route(anyCollection(), eq(dataFlowComponents), any(RouterGetRecordStrategy.class), any(BiConsumer.class));
@@ -451,7 +455,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_InactiveAcknowledge
RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2);
Record rec = records.get(0);
event = mock(JacksonEvent.class);
- eventHandle = mock(EventHandle.class);
+ eventHandle = mock(DefaultEventHandle.class);
when(event.getEventHandle()).thenReturn(null);
when(rec.getData()).thenReturn(event);
routerCopyRecordStrategy.getRecord(rec);
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java
index a9c030793b..4c56113323 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java
@@ -42,10 +42,9 @@
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.event.EventFactory;
-import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.event.EventBuilder;
import org.opensearch.dataprepper.model.event.EventMetadata;
-import org.opensearch.dataprepper.event.DefaultEventHandle;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
@@ -67,7 +66,7 @@ public class RouterCopyRecordStrategyTests {
private JacksonEvent event;
- private Map handleRefCount;
+ private Map handleRefCount;
private static class TestComponent {
}
@@ -80,11 +79,11 @@ void setUp() {
acknowledgementSet1 = mock(AcknowledgementSet.class);
try {
lenient().doAnswer((i) -> {
- EventHandle handle = (EventHandle) i.getArgument(0);
+ DefaultEventHandle handle = (DefaultEventHandle) i.getArgument(0);
int v = handleRefCount.getOrDefault(handle, 0);
handleRefCount.put(handle, v+1);
return null;
- }).when(acknowledgementSetManager).acquireEventReference(any(EventHandle.class));
+ }).when(acknowledgementSetManager).acquireEventReference(any(DefaultEventHandle.class));
} catch (Exception e){}
mockRecordsIn = IntStream.range(0, 10)
.mapToObj(i -> mock(Record.class))
@@ -99,12 +98,12 @@ void setUp() {
.collect(Collectors.toList());
}
- private void attachEventHandlesToRecordsIn(List eventHandles) {
+ private void attachEventHandlesToRecordsIn(List eventHandles) {
Iterator iter = recordsIn.iterator();
while (iter.hasNext()) {
Record r = (Record) iter.next();
- DefaultEventHandle handle = new DefaultEventHandle(acknowledgementSet1);
- ((JacksonEvent)r.getData()).setEventHandle(handle);
+ DefaultEventHandle handle = (DefaultEventHandle)((JacksonEvent)r.getData()).getEventHandle();
+ handle.setAcknowledgementSet(acknowledgementSet1);
eventHandles.add(handle);
}
}
@@ -187,10 +186,10 @@ void test_one_record_with_acknowledgements() {
= Collections.singletonList(dataFlowComponent);
final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents);
- List eventHandles = new ArrayList<>();
+ List eventHandles = new ArrayList<>();
attachEventHandlesToRecordsIn(eventHandles);
Record firstRecord = recordsIn.iterator().next();
- EventHandle firstHandle = ((Event)firstRecord.getData()).getEventHandle();
+ DefaultEventHandle firstHandle = (DefaultEventHandle)((Event)firstRecord.getData()).getEventHandle();
Record recordOut = getRecordStrategy.getRecord(firstRecord);
assertThat(recordOut, sameInstance(firstRecord));
assertTrue(getRecordStrategy.getReferencedRecords().contains(firstRecord));
@@ -209,7 +208,7 @@ void test_multiple_records_with_acknowledgements() {
= Collections.singletonList(dataFlowComponent);
final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents);
- List eventHandles = new ArrayList<>();
+ List eventHandles = new ArrayList<>();
attachEventHandlesToRecordsIn(eventHandles);
Collection recordsOut = getRecordStrategy.getAllRecords(recordsIn);
assertThat(recordsOut.size(), equalTo(recordsIn.size()));
@@ -238,12 +237,12 @@ void test_one_record_with_acknowledgements_and_multi_components() {
}
final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents);
- List eventHandles = new ArrayList<>();
+ List eventHandles = new ArrayList<>();
attachEventHandlesToRecordsIn(eventHandles);
try {
doAnswer((i) -> {
JacksonEvent e1 = (JacksonEvent) i.getArgument(0);
- e1.setEventHandle(new DefaultEventHandle(acknowledgementSet1));
+ ((DefaultEventHandle)e1.getEventHandle()).setAcknowledgementSet(acknowledgementSet1);
return null;
}).when(acknowledgementSet1).add(any(JacksonEvent.class));
} catch (Exception e){}
@@ -254,14 +253,14 @@ void test_one_record_with_acknowledgements_and_multi_components() {
when(eventBuilder.build()).thenReturn(JacksonEvent.fromEvent(event));
when(eventFactory.eventBuilder(EventBuilder.class)).thenReturn(eventBuilder);
Record firstRecord = recordsIn.iterator().next();
- EventHandle firstHandle = ((Event)firstRecord.getData()).getEventHandle();
+ DefaultEventHandle firstHandle = (DefaultEventHandle)((Event)firstRecord.getData()).getEventHandle();
Record recordOut = getRecordStrategy.getRecord(firstRecord);
assertThat(recordOut, sameInstance(firstRecord));
assertTrue(getRecordStrategy.getReferencedRecords().contains(firstRecord));
recordOut = getRecordStrategy.getRecord(firstRecord);
assertThat(recordOut, not(sameInstance(firstRecord)));
assertFalse(handleRefCount.containsKey(firstHandle));
- EventHandle newHandle = ((JacksonEvent)recordOut.getData()).getEventHandle();
+ DefaultEventHandle newHandle = (DefaultEventHandle)((JacksonEvent)recordOut.getData()).getEventHandle();
assertTrue(getRecordStrategy.getReferencedRecords().contains(recordOut));
assertThat(newHandle, not(equalTo(null)));
assertFalse(handleRefCount.containsKey(newHandle));
@@ -276,12 +275,12 @@ void test_multiple_records_with_acknowledgements_and_multi_components() {
}
final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents);
- List eventHandles = new ArrayList<>();
+ List eventHandles = new ArrayList<>();
attachEventHandlesToRecordsIn(eventHandles);
try {
doAnswer((i) -> {
JacksonEvent e1 = (JacksonEvent) i.getArgument(0);
- e1.setEventHandle(new DefaultEventHandle(acknowledgementSet1));
+ ((DefaultEventHandle)e1.getEventHandle()).setAcknowledgementSet(acknowledgementSet1);
return null;
}).when(acknowledgementSet1).add(any(JacksonEvent.class));
} catch (Exception e){}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java
index a43b9c17c4..1c920b6fc3 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java
@@ -7,6 +7,7 @@
import org.mockito.ArgumentCaptor;
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
+import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
import org.opensearch.dataprepper.model.configuration.PluginSetting;
import org.opensearch.dataprepper.model.plugin.NoPluginFoundException;
@@ -57,6 +58,7 @@ class DefaultPluginFactoryTest {
private BeanFactory beanFactory;
private String pipelineName;
private DefaultAcknowledgementSetManager acknowledgementSetManager;
+ private CircuitBreaker circuitBreaker;
private DefaultEventFactory eventFactory;
private PluginConfigurationObservableFactory pluginConfigurationObservableFactory;
private PluginConfigObservable pluginConfigObservable;
@@ -68,6 +70,7 @@ void setUp() {
pluginConfigurationConverter = mock(PluginConfigurationConverter.class);
acknowledgementSetManager = mock(DefaultAcknowledgementSetManager.class);
+ circuitBreaker = mock(CircuitBreaker.class);
eventFactory = mock(DefaultEventFactory.class);
pluginProviders = new ArrayList<>();
given(pluginProviderLoader.getPluginProviders()).willReturn(pluginProviders);
@@ -94,7 +97,7 @@ void setUp() {
private DefaultPluginFactory createObjectUnderTest() {
return new DefaultPluginFactory(
pluginProviderLoader, pluginCreator, pluginConfigurationConverter, beanFactoryProvider, eventFactory,
- acknowledgementSetManager, pluginConfigurationObservableFactory);
+ acknowledgementSetManager, pluginConfigurationObservableFactory, circuitBreaker);
}
@Test
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java
new file mode 100644
index 0000000000..5075c210ac
--- /dev/null
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.plugins;
+
+import org.opensearch.dataprepper.model.CheckpointState;
+import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
+import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.record.Record;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.TimeoutException;
+
+@DataPrepperPlugin(name = "test_off_heap", pluginType = Buffer.class)
+public class TestOffHeapBuffer implements Buffer {
+ @Override
+ public void write(Record record, int timeoutInMillis) throws TimeoutException {
+
+ }
+
+ @Override
+ public void writeAll(Collection records, int timeoutInMillis) throws Exception {
+
+ }
+
+ @Override
+ public Map.Entry read(int timeoutInMillis) {
+ return null;
+ }
+
+ @Override
+ public void checkpoint(CheckpointState checkpointState) {
+
+ }
+
+ @Override
+ public boolean isWrittenOffHeapOnly() {
+ return true;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return false;
+ }
+}
diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java
index 4e6e40c57a..6a1122ee56 100644
--- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java
+++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java
@@ -163,7 +163,7 @@ void test_saveProgressStateForPartition() {
Optional sourcePartition = coordinator.acquireAvailablePartition(DEFAULT_PARTITION_TYPE);
assertThat(sourcePartition.isPresent(), equalTo(true));
TestEnhancedSourcePartition partition = (TestEnhancedSourcePartition) sourcePartition.get();
- coordinator.saveProgressStateForPartition(partition);
+ coordinator.saveProgressStateForPartition(partition, null);
verify(sourceCoordinationStore).tryAcquireAvailablePartition(anyString(), anyString(), any(Duration.class));
verify(sourceCoordinationStore).tryUpdateSourcePartitionItem(any(SourcePartitionStoreItem.class));
diff --git a/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml b/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml
new file mode 100644
index 0000000000..e7ed4c35e1
--- /dev/null
+++ b/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml
@@ -0,0 +1,7 @@
+test-pipeline-1:
+ source:
+ stdin:
+ buffer:
+ test_off_heap:
+ sink:
+ - stdout:
diff --git a/data-prepper-expression/build.gradle b/data-prepper-expression/build.gradle
index 41a5a44263..70d3c1d78d 100644
--- a/data-prepper-expression/build.gradle
+++ b/data-prepper-expression/build.gradle
@@ -24,7 +24,7 @@ dependencies {
implementation(libs.spring.context) {
exclude group: 'commons-logging', module: 'commons-logging'
}
- implementation platform('org.apache.logging.log4j:log4j-bom:2.20.0')
+ implementation platform('org.apache.logging.log4j:log4j-bom:2.21.1')
implementation 'org.apache.logging.log4j:log4j-core'
implementation 'org.apache.logging.log4j:log4j-slf4j2-impl'
implementation 'com.github.seancfoley:ipaddress:5.4.0'
diff --git a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java
index 072b3a7393..6653e5c7b4 100644
--- a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java
+++ b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java
@@ -7,6 +7,7 @@
import org.antlr.v4.runtime.tree.ParseTree;
import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.event.JacksonEvent;
import javax.inject.Inject;
import javax.inject.Named;
@@ -52,4 +53,25 @@ public Boolean isValidExpressionStatement(final String statement) {
return false;
}
}
+
+ @Override
+ public Boolean isValidFormatExpression(final String format) {
+ int fromIndex = 0;
+ int position = 0;
+ while ((position = format.indexOf("${", fromIndex)) != -1) {
+ int endPosition = format.indexOf("}", position + 1);
+ if (endPosition == -1) {
+ return false;
+ }
+ String name = format.substring(position + 2, endPosition);
+
+ Object val;
+ // Invalid if it is not a valid key and not a valid expression statement
+ if (!JacksonEvent.isValidEventKey(name) && !isValidExpressionStatement(name)) {
+ return false;
+ }
+ fromIndex = endPosition + 1;
+ }
+ return true;
+ }
}
diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java
index d68808dc85..a91e7fe368 100644
--- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java
+++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java
@@ -5,22 +5,26 @@
package org.opensearch.dataprepper.expression;
-import org.opensearch.dataprepper.model.event.Event;
import org.antlr.v4.runtime.tree.ParseTree;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.dataprepper.model.event.Event;
-import java.util.UUID;
import java.util.Random;
+import java.util.UUID;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
@@ -119,5 +123,24 @@ void isValidExpressionStatement_returns_false_when_parse_throws() {
assertThat(result, equalTo(false));
}
+ @ParameterizedTest
+ @CsvSource({
+ "abc-${/foo, false",
+ "abc-${/foo}, true",
+ "abc-${getMetadata(\"key\")}, true",
+ "abc-${getXYZ(\"key\")}, true",
+ "abc-${invalid, false"
+ })
+ void isValidFormatExpressionsReturnsCorrectResult(final String format, final Boolean expectedResult) {
+ assertThat(statementEvaluator.isValidFormatExpression(format), equalTo(expectedResult));
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"abc-${anyS(=tring}"})
+ void isValidFormatExpressionsReturnsFalseWhenIsValidKeyAndValidExpressionIsFalse(final String format) {
+ doThrow(RuntimeException.class).when(parser).parse(anyString());
+ assertThat(statementEvaluator.isValidFormatExpression(format), equalTo(false));
+ }
+
}
diff --git a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java
index a27d9ea4cb..fc9963ab46 100644
--- a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java
+++ b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java
@@ -104,9 +104,7 @@ private void stageLogEvents() {
}
private void addToBuffer(final Record log, final String logString) {
- if (log.getData().getEventHandle() != null) {
- bufferedEventHandles.add(log.getData().getEventHandle());
- }
+ bufferedEventHandles.add(log.getData().getEventHandle());
buffer.writeEvent(logString.getBytes(StandardCharsets.UTF_8));
}
}
diff --git a/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java b/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java
index e8e416d53e..8e7fcebc4e 100644
--- a/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java
+++ b/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java
@@ -8,7 +8,6 @@
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.opensearch.dataprepper.model.event.Event;
-import org.opensearch.dataprepper.model.event.EventHandle;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.sink.cloudwatch_logs.buffer.Buffer;
@@ -62,8 +61,6 @@ Collection> getSampleRecordsCollectionSmall() {
final ArrayList> returnCollection = new ArrayList<>();
for (int i = 0; i < 5; i++) {
JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage");
- final EventHandle mockEventHandle = mock(EventHandle.class);
- mockJacksonEvent.setEventHandle(mockEventHandle);
returnCollection.add(new Record<>(mockJacksonEvent));
}
@@ -74,8 +71,6 @@ Collection> getSampleRecordsCollection() {
final ArrayList> returnCollection = new ArrayList<>();
for (int i = 0; i < thresholdConfig.getBatchSize(); i++) {
JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage");
- final EventHandle mockEventHandle = mock(EventHandle.class);
- mockJacksonEvent.setEventHandle(mockEventHandle);
returnCollection.add(new Record<>(mockJacksonEvent));
}
@@ -86,8 +81,6 @@ Collection> getSampleRecordsOfLargerSize() {
final ArrayList> returnCollection = new ArrayList<>();
for (int i = 0; i < thresholdConfig.getBatchSize() * 2; i++) {
JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("a".repeat((int) (thresholdConfig.getMaxRequestSizeBytes()/24)));
- final EventHandle mockEventHandle = mock(EventHandle.class);
- mockJacksonEvent.setEventHandle(mockEventHandle);
returnCollection.add(new Record<>(mockJacksonEvent));
}
@@ -98,8 +91,6 @@ Collection> getSampleRecordsOfLimitSize() {
final ArrayList> returnCollection = new ArrayList<>();
for (int i = 0; i < thresholdConfig.getBatchSize(); i++) {
JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage".repeat((int) thresholdConfig.getMaxEventSizeBytes()));
- final EventHandle mockEventHandle = mock(EventHandle.class);
- mockJacksonEvent.setEventHandle(mockEventHandle);
returnCollection.add(new Record<>(mockJacksonEvent));
}
diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java
index c6f6018a4f..29dd0c204d 100644
--- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java
+++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java
@@ -6,6 +6,7 @@
package org.opensearch.dataprepper.plugins.accumulator;
import java.io.IOException;
+import java.io.OutputStream;
/**
* A buffer can hold data before flushing it any Sink.
@@ -21,5 +22,8 @@ public interface Buffer {
long getDuration();
byte[] getSinkBufferData() throws IOException;
- void writeEvent(byte[] bytes) throws IOException;
+
+ OutputStream getOutputStream();
+
+ void setEventCount(int eventCount);
}
diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java
index e583e54260..252dac88a9 100644
--- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java
+++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java
@@ -8,6 +8,7 @@
import org.apache.commons.lang3.time.StopWatch;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
+import java.io.OutputStream;
import java.util.concurrent.TimeUnit;
/**
@@ -49,16 +50,13 @@ public byte[] getSinkBufferData() throws IOException {
return byteArrayOutputStream.toByteArray();
}
- /**
- * write byte array to output stream.
- *
- * @param bytes byte array.
- * @throws IOException while writing to output stream fails.
- */
@Override
- public void writeEvent(byte[] bytes) throws IOException {
- byteArrayOutputStream.write(bytes);
- byteArrayOutputStream.write(System.lineSeparator().getBytes());
- eventCount++;
+ public OutputStream getOutputStream() {
+ return byteArrayOutputStream;
+ }
+
+ @Override
+ public void setEventCount(int eventCount) {
+ this.eventCount = eventCount;
}
}
\ No newline at end of file
diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java
index 9f9b4a3aac..f419fde25b 100644
--- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java
+++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java
@@ -70,18 +70,6 @@ public byte[] getSinkBufferData() throws IOException {
return fileData;
}
- /**
- * write byte array to output stream.
- * @param bytes byte array.
- * @throws IOException while writing to output stream fails.
- */
- @Override
- public void writeEvent(byte[] bytes) throws IOException {
- outputStream.write(bytes);
- outputStream.write(System.lineSeparator().getBytes());
- eventCount++;
- }
-
/**
* Flushing the buffered data into the output stream.
*/
@@ -106,4 +94,14 @@ protected void removeTemporaryFile() {
}
}
}
+
+ @Override
+ public OutputStream getOutputStream() {
+ return outputStream;
+ }
+
+ @Override
+ public void setEventCount(int eventCount) {
+ this.eventCount = eventCount;
+ }
}
\ No newline at end of file
diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java
index ad07cc4011..2222919ed0 100644
--- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java
+++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java
@@ -10,12 +10,12 @@
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.junit.jupiter.MockitoExtension;
-import java.io.IOException;
-
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.io.IOException;
@ExtendWith(MockitoExtension.class)
class InMemoryBufferTest {
@@ -28,8 +28,10 @@ class InMemoryBufferTest {
void test_with_write_event_into_buffer() throws IOException {
inMemoryBuffer = new InMemoryBuffer();
+ int i = 0;
while (inMemoryBuffer.getEventCount() < MAX_EVENTS) {
- inMemoryBuffer.writeEvent(generateByteArray());
+ inMemoryBuffer.getOutputStream().write(generateByteArray());
+ inMemoryBuffer.setEventCount(++i);
}
assertThat(inMemoryBuffer.getSize(), greaterThanOrEqualTo(54110L));
assertThat(inMemoryBuffer.getEventCount(), equalTo(MAX_EVENTS));
@@ -53,4 +55,4 @@ private byte[] generateByteArray() {
}
return bytes;
}
-}
\ No newline at end of file
+}
diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java
index 53c556e75c..e6d65a23aa 100644
--- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java
+++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java
@@ -24,6 +24,7 @@
@ExtendWith(MockitoExtension.class)
class LocalFileBufferTest {
+ public static final int MAX_EVENTS = 55;
public static final String KEY = UUID.randomUUID().toString() + ".log";
public static final String PREFIX = "local";
@@ -40,8 +41,10 @@ void setUp() throws IOException {
@Test
void test_with_write_events_into_buffer() throws IOException {
- while (localFileBuffer.getEventCount() < 55) {
- localFileBuffer.writeEvent(generateByteArray());
+ int i = 0;
+ while (localFileBuffer.getEventCount() < MAX_EVENTS) {
+ localFileBuffer.getOutputStream().write(generateByteArray());
+ localFileBuffer.setEventCount(++i);
}
assertThat(localFileBuffer.getSize(), greaterThan(1l));
assertThat(localFileBuffer.getEventCount(), equalTo(55));
@@ -81,4 +84,4 @@ private byte[] generateByteArray() {
}
return bytes;
}
-}
\ No newline at end of file
+}
diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java
index d21a664fb5..20667e1564 100644
--- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java
+++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java
@@ -14,7 +14,6 @@
import java.io.IOException;
import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
class ThresholdValidatorTest {
@@ -23,20 +22,6 @@ class ThresholdValidatorTest {
@BeforeEach
void setUp() throws IOException {
inMemoryBuffer = new InMemoryBufferFactory().getBuffer();
-
- while (inMemoryBuffer.getEventCount() < 100) {
- inMemoryBuffer.writeEvent(generateByteArray());
- }
- }
-
- @Test
- void test_exceedThreshold_true_dueTo_maxEvents_is_less_than_buffered_event_count() {
- final int maxEvents = 95;
- final ByteCount maxBytes = ByteCount.parse("50kb");
- final long maxCollectionDuration = 15;
- boolean isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents,
- maxBytes, maxCollectionDuration);
- assertTrue(isThresholdExceed, "Threshold not exceeded");
}
@Test
@@ -49,16 +34,6 @@ void test_exceedThreshold_false_dueTo_maxEvents_is_greater_than_buffered_event_c
assertFalse(isThresholdExceed, "Threshold exceeded");
}
- @Test
- void test_exceedThreshold_ture_dueTo_maxBytes_is_less_than_buffered_byte_count() {
- final int maxEvents = 500;
- final ByteCount maxBytes = ByteCount.parse("1b");
- final long maxCollectionDuration = 15;
- boolean isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents, maxBytes,
- maxCollectionDuration);
- assertTrue(isThresholdExceed, "Threshold not exceeded");
- }
-
@Test
void test_exceedThreshold_false_dueTo_maxBytes_is_greater_than_buffered_byte_count() {
final int maxEvents = 500;
@@ -69,58 +44,4 @@ void test_exceedThreshold_false_dueTo_maxBytes_is_greater_than_buffered_byte_cou
assertFalse(isThresholdExceed, "Threshold exceeded");
}
- @Test
- void test_exceedThreshold_ture_dueTo_maxCollectionDuration_is_less_than_buffered_event_collection_duration()
- throws IOException, InterruptedException {
- final int maxEvents = 500;
- final ByteCount maxBytes = ByteCount.parse("500mb");
- final long maxCollectionDuration = 10;
-
- inMemoryBuffer = new InMemoryBufferFactory().getBuffer();
- boolean isThresholdExceed = Boolean.FALSE;
- synchronized (this) {
- while (inMemoryBuffer.getEventCount() < 100) {
- inMemoryBuffer.writeEvent(generateByteArray());
- isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents,
- maxBytes, maxCollectionDuration);
- if (isThresholdExceed) {
- break;
- }
- wait(5000);
- }
- }
- assertTrue(isThresholdExceed, "Threshold not exceeded");
- }
-
- @Test
- void test_exceedThreshold_ture_dueTo_maxCollectionDuration_is_greater_than_buffered_event_collection_duration()
- throws IOException, InterruptedException {
- final int maxEvents = 500;
- final ByteCount maxBytes = ByteCount.parse("500mb");
- final long maxCollectionDuration = 240;
-
- inMemoryBuffer = new InMemoryBufferFactory().getBuffer();
-
- boolean isThresholdExceed = Boolean.FALSE;
- synchronized (this) {
- while (inMemoryBuffer.getEventCount() < 100) {
- inMemoryBuffer.writeEvent(generateByteArray());
- isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer,
- maxEvents, maxBytes, maxCollectionDuration);
- if (isThresholdExceed) {
- break;
- }
- wait(50);
- }
- }
- assertFalse(isThresholdExceed, "Threshold exceeded");
- }
-
- private byte[] generateByteArray() {
- byte[] bytes = new byte[10000];
- for (int i = 0; i < 10000; i++) {
- bytes[i] = (byte) i;
- }
- return bytes;
- }
}
diff --git a/data-prepper-plugins/date-processor/README.md b/data-prepper-plugins/date-processor/README.md
index a5c147ecd6..c4d9acd4f5 100644
--- a/data-prepper-plugins/date-processor/README.md
+++ b/data-prepper-plugins/date-processor/README.md
@@ -104,6 +104,8 @@ processor:
* Type: String
* Default: `Locale.ROOT`
+* `to_origination_metadata` (Optional): When this option is used, matched time is put into the event's metadata as an instance of `Instant`.
+
## Metrics
* `dateProcessingMatchSuccessCounter`: Number of records that match with at least one pattern specified in match configuration option.
diff --git a/data-prepper-plugins/date-processor/build.gradle b/data-prepper-plugins/date-processor/build.gradle
index 6433af9b6d..743761201d 100644
--- a/data-prepper-plugins/date-processor/build.gradle
+++ b/data-prepper-plugins/date-processor/build.gradle
@@ -12,5 +12,6 @@ dependencies {
implementation project(':data-prepper-test-common')
implementation 'com.fasterxml.jackson.core:jackson-databind'
implementation 'io.micrometer:micrometer-core'
+ implementation libs.commons.lang3
testImplementation libs.commons.lang3
}
diff --git a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java
index b9094926ee..d0808053b2 100644
--- a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java
+++ b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java
@@ -27,6 +27,7 @@
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
+import org.apache.commons.lang3.tuple.Pair;
@DataPrepperPlugin(name = "date", pluginType = Processor.class, pluginConfigurationType = DateProcessorConfig.class)
public class DateProcessor extends AbstractProcessor, Record> {
@@ -71,7 +72,16 @@ public Collection> doExecute(Collection> records) {
zonedDateTime = getDateTimeFromTimeReceived(record);
else if (keyToParse != null && !keyToParse.isEmpty()) {
- zonedDateTime = getDateTimeFromMatch(record);
+ Pair result = getDateTimeFromMatch(record);
+ if (result != null) {
+ zonedDateTime = result.getLeft();
+ Instant timeStamp = result.getRight();
+ if (dateProcessorConfig.getToOriginationMetadata()) {
+ Event event = (Event)record.getData();
+ event.getMetadata().setExternalOriginationTime(timeStamp);
+ event.getEventHandle().setExternalOriginationTime(timeStamp);
+ }
+ }
populateDateProcessorMetrics(zonedDateTime);
}
@@ -119,7 +129,7 @@ private String getDateTimeFromTimeReceived(final Record record) {
return timeReceived.atZone(dateProcessorConfig.getDestinationZoneId()).format(getOutputFormatter());
}
- private String getDateTimeFromMatch(final Record record) {
+ private Pair getDateTimeFromMatch(final Record record) {
final String sourceTimestamp = getSourceTimestamp(record);
if (sourceTimestamp == null)
return null;
@@ -136,12 +146,12 @@ private String getSourceTimestamp(final Record record) {
}
}
- private String getFormattedDateTimeString(final String sourceTimestamp) {
+ private Pair getFormattedDateTimeString(final String sourceTimestamp) {
for (DateTimeFormatter formatter : dateTimeFormatters) {
try {
- return ZonedDateTime.parse(sourceTimestamp, formatter).format(getOutputFormatter().withZone(dateProcessorConfig.getDestinationZoneId()));
+ ZonedDateTime tmp = ZonedDateTime.parse(sourceTimestamp, formatter);
+ return Pair.of(tmp.format(getOutputFormatter().withZone(dateProcessorConfig.getDestinationZoneId())), tmp.toInstant());
} catch (Exception ignored) {
-
}
}
diff --git a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java
index fd6c1f25be..5e06b48cbb 100644
--- a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java
+++ b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java
@@ -15,6 +15,7 @@
public class DateProcessorConfig {
static final Boolean DEFAULT_FROM_TIME_RECEIVED = false;
+ static final Boolean DEFAULT_TO_ORIGINATION_METADATA = false;
static final String DEFAULT_DESTINATION = "@timestamp";
static final String DEFAULT_SOURCE_TIMEZONE = ZoneId.systemDefault().toString();
static final String DEFAULT_DESTINATION_TIMEZONE = ZoneId.systemDefault().toString();
@@ -45,6 +46,9 @@ public List getPatterns() {
@JsonProperty("from_time_received")
private Boolean fromTimeReceived = DEFAULT_FROM_TIME_RECEIVED;
+ @JsonProperty("to_origination_metadata")
+ private Boolean toOriginationMetadata = DEFAULT_TO_ORIGINATION_METADATA;
+
@JsonProperty("match")
private List match;
@@ -76,6 +80,10 @@ public Boolean getFromTimeReceived() {
return fromTimeReceived;
}
+ public Boolean getToOriginationMetadata() {
+ return toOriginationMetadata;
+ }
+
public List getMatch() {
return match;
}
diff --git a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java
index 0959c9db2f..db604039fa 100644
--- a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java
+++ b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java
@@ -51,8 +51,9 @@ void isValidMatchAndFromTimestampReceived_should_return_true_if_from_time_receiv
}
@Test
- void isValidMatchAndFromTimestampReceived_should_return_false_if_from_time_received_and_match_are_not_configured() {
- assertThat(dateProcessorConfig.isValidMatchAndFromTimestampReceived(), equalTo(false));
+ void testToOriginationMetadata_should_return_true() throws NoSuchFieldException, IllegalAccessException {
+ reflectivelySetField(dateProcessorConfig, "toOriginationMetadata", true);
+ assertThat(dateProcessorConfig.getToOriginationMetadata(), equalTo(true));
}
@Test
@@ -178,4 +179,4 @@ private void reflectivelySetField(final DateProcessorConfig dateProcessorConfig,
field.setAccessible(false);
}
}
-}
\ No newline at end of file
+}
diff --git a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java
index a45daeb56b..ce3906a635 100644
--- a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java
+++ b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java
@@ -83,7 +83,8 @@ void setup() {
lenient().when(pluginMetrics.counter(DateProcessor.DATE_PROCESSING_MATCH_SUCCESS)).thenReturn(dateProcessingMatchSuccessCounter);
lenient().when(pluginMetrics.counter(DateProcessor.DATE_PROCESSING_MATCH_FAILURE)).thenReturn(dateProcessingMatchFailureCounter);
when(mockDateProcessorConfig.getDateWhen()).thenReturn(null);
- expectedDateTime = LocalDateTime.now();
+ expectedInstant = Instant.now();
+ expectedDateTime = LocalDateTime.ofInstant(expectedInstant, ZoneId.systemDefault());
}
@AfterEach
@@ -361,6 +362,35 @@ void match_with_different_year_formats_test(String pattern) {
verify(dateProcessingMatchSuccessCounter, times(1)).increment();
}
+ @ParameterizedTest
+ @ValueSource(strings = {"yyyy MM dd HH mm ss"})
+ void match_with_to_origination_metadata(String pattern) {
+ when(mockDateMatch.getKey()).thenReturn("logDate");
+ when(mockDateMatch.getPatterns()).thenReturn(Collections.singletonList(pattern));
+
+ List dateMatches = Collections.singletonList(mockDateMatch);
+ when(mockDateProcessorConfig.getMatch()).thenReturn(dateMatches);
+ when(mockDateProcessorConfig.getSourceZoneId()).thenReturn(ZoneId.systemDefault());
+ when(mockDateProcessorConfig.getDestinationZoneId()).thenReturn(ZoneId.systemDefault());
+ when(mockDateProcessorConfig.getSourceLocale()).thenReturn(Locale.ROOT);
+ when(mockDateProcessorConfig.getToOriginationMetadata()).thenReturn(true);
+
+ dateProcessor = createObjectUnderTest();
+
+ Map testData = getTestData();
+ testData.put("logDate", expectedDateTime.format(DateTimeFormatter.ofPattern(pattern)));
+
+ final Record record = buildRecordWithEvent(testData);
+ final List> processedRecords = (List>) dateProcessor.doExecute(Collections.singletonList(record));
+
+ Event event = (Event)processedRecords.get(0).getData();
+ Assertions.assertTrue(event.getMetadata().getExternalOriginationTime() != null);
+ Assertions.assertTrue(event.getEventHandle().getExternalOriginationTime() != null);
+ ZonedDateTime expectedZonedDatetime = expectedDateTime.atZone(mockDateProcessorConfig.getSourceZoneId()).truncatedTo(ChronoUnit.SECONDS);
+ Assertions.assertTrue(expectedZonedDatetime.equals(event.getMetadata().getExternalOriginationTime().atZone(mockDateProcessorConfig.getSourceZoneId())));
+ verify(dateProcessingMatchSuccessCounter, times(1)).increment();
+ }
+
@ParameterizedTest
@ValueSource(strings = {"MMM/dd", "MM dd"})
void match_without_year_test(String pattern) {
diff --git a/data-prepper-plugins/dynamodb-source/README.md b/data-prepper-plugins/dynamodb-source/README.md
index 1585e66e12..db620a6406 100644
--- a/data-prepper-plugins/dynamodb-source/README.md
+++ b/data-prepper-plugins/dynamodb-source/README.md
@@ -52,6 +52,8 @@ source:
* s3_bucket (Required): The destination bucket to store the exported data files
* s3_prefix (Optional): Custom prefix.
+* s3_sse_kms_key_id (Optional): A AWS KMS Customer Managed Key (CMK) to encrypt the export data files. The key id will
+ be the ARN of the Key, e.g. arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad4-80ca-63b12b3ec147
### Stream Configurations
diff --git a/data-prepper-plugins/dynamodb-source/build.gradle b/data-prepper-plugins/dynamodb-source/build.gradle
index 71b8dc2afc..1fc14af6ea 100644
--- a/data-prepper-plugins/dynamodb-source/build.gradle
+++ b/data-prepper-plugins/dynamodb-source/build.gradle
@@ -26,10 +26,12 @@ dependencies {
implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-ion'
implementation project(path: ':data-prepper-plugins:aws-plugin-api')
+ implementation project(path: ':data-prepper-plugins:buffer-common')
testImplementation platform('org.junit:junit-bom:5.9.1')
testImplementation 'org.junit.jupiter:junit-jupiter'
+ testImplementation testLibs.mockito.inline
testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java
index 7d72049a6a..0704694d0c 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java
@@ -8,7 +8,10 @@
import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions;
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.AwsAuthenticationConfig;
+import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.ExportConfig;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient;
import software.amazon.awssdk.services.s3.S3Client;
@@ -17,8 +20,11 @@ public class ClientFactory {
private final AwsCredentialsProvider awsCredentialsProvider;
private final AwsAuthenticationConfig awsAuthenticationConfig;
+ private final ExportConfig exportConfig;
- public ClientFactory(AwsCredentialsSupplier awsCredentialsSupplier, AwsAuthenticationConfig awsAuthenticationConfig) {
+ public ClientFactory(final AwsCredentialsSupplier awsCredentialsSupplier,
+ final AwsAuthenticationConfig awsAuthenticationConfig,
+ final ExportConfig exportConfig) {
awsCredentialsProvider = awsCredentialsSupplier.getProvider(AwsCredentialsOptions.builder()
.withRegion(awsAuthenticationConfig.getAwsRegion())
.withStsRoleArn(awsAuthenticationConfig.getAwsStsRoleArn())
@@ -26,6 +32,7 @@ public ClientFactory(AwsCredentialsSupplier awsCredentialsSupplier, AwsAuthentic
.withStsHeaderOverrides(awsAuthenticationConfig.getAwsStsHeaderOverrides())
.build());
this.awsAuthenticationConfig = awsAuthenticationConfig;
+ this.exportConfig = exportConfig;
}
@@ -47,8 +54,20 @@ public DynamoDbClient buildDynamoDBClient() {
public S3Client buildS3Client() {
return S3Client.builder()
+ .region(getS3ClientRegion())
.credentialsProvider(awsCredentialsProvider)
+ .overrideConfiguration(ClientOverrideConfiguration.builder()
+ .retryPolicy(retryPolicy -> retryPolicy.numRetries(5).build())
+ .build())
.build();
}
+ private Region getS3ClientRegion() {
+ if (exportConfig != null && exportConfig.getAwsRegion() != null) {
+ return exportConfig.getAwsRegion();
+ }
+
+ return awsAuthenticationConfig.getAwsRegion();
+ }
+
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java
index 6b5430d997..fed9e83f79 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java
@@ -6,6 +6,7 @@
package org.opensearch.dataprepper.plugins.source.dynamodb;
import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException;
@@ -41,6 +42,7 @@
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient;
import software.amazon.awssdk.services.s3.S3Client;
+import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.Map;
@@ -58,6 +60,8 @@ public class DynamoDBService {
private final EnhancedSourceCoordinator coordinator;
private final DynamoDbClient dynamoDbClient;
+
+ private final DynamoDBSourceConfig dynamoDBSourceConfig;
//
private final DynamoDbStreamsClient dynamoDbStreamsClient;
@@ -69,10 +73,21 @@ public class DynamoDBService {
private final PluginMetrics pluginMetrics;
+ private final AcknowledgementSetManager acknowledgementSetManager;
+
+ static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60);
+ static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000;
+
- public DynamoDBService(EnhancedSourceCoordinator coordinator, ClientFactory clientFactory, DynamoDBSourceConfig sourceConfig, PluginMetrics pluginMetrics) {
+ public DynamoDBService(final EnhancedSourceCoordinator coordinator,
+ final ClientFactory clientFactory,
+ final DynamoDBSourceConfig sourceConfig,
+ final PluginMetrics pluginMetrics,
+ final AcknowledgementSetManager acknowledgementSetManager) {
this.coordinator = coordinator;
this.pluginMetrics = pluginMetrics;
+ this.acknowledgementSetManager = acknowledgementSetManager;
+ this.dynamoDBSourceConfig = sourceConfig;
// Initialize AWS clients
dynamoDbClient = clientFactory.buildDynamoDBClient();
@@ -99,10 +114,10 @@ public void start(Buffer> buffer) {
Runnable exportScheduler = new ExportScheduler(coordinator, dynamoDbClient, manifestFileReader, pluginMetrics);
DataFileLoaderFactory loaderFactory = new DataFileLoaderFactory(coordinator, s3Client, pluginMetrics, buffer);
- Runnable fileLoaderScheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics);
+ Runnable fileLoaderScheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig);
ShardConsumerFactory consumerFactory = new ShardConsumerFactory(coordinator, dynamoDbStreamsClient, pluginMetrics, shardManager, buffer);
- Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics);
+ Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig);
// May consider start or shutdown the scheduler on demand
// Currently, event after the exports are done, the related scheduler will not be shutdown
@@ -156,14 +171,18 @@ public void init() {
Instant startTime = Instant.now();
if (tableInfo.getMetadata().isExportRequired()) {
-// exportTime = Instant.now();
- createExportPartition(tableInfo.getTableArn(), startTime, tableInfo.getMetadata().getExportBucket(), tableInfo.getMetadata().getExportPrefix());
+ createExportPartition(
+ tableInfo.getTableArn(),
+ startTime,
+ tableInfo.getMetadata().getExportBucket(),
+ tableInfo.getMetadata().getExportPrefix(),
+ tableInfo.getMetadata().getExportKmsKeyId());
}
if (tableInfo.getMetadata().isStreamRequired()) {
List shardIds;
// start position by default is TRIM_HORIZON if not provided.
- if (tableInfo.getMetadata().isExportRequired() || String.valueOf(StreamStartPosition.LATEST).equals(tableInfo.getMetadata().getStreamStartPosition())) {
+ if (tableInfo.getMetadata().isExportRequired() || tableInfo.getMetadata().getStreamStartPosition() == StreamStartPosition.LATEST) {
// For a continued data extraction process that involves both export and stream
// The export must be completed and loaded before stream can start.
// Moreover, there should not be any gaps between the export time and the time start reading the stream
@@ -195,11 +214,12 @@ public void init() {
* @param bucket Export bucket
* @param prefix Export Prefix
*/
- private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix) {
+ private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix, String kmsKeyId) {
ExportProgressState exportProgressState = new ExportProgressState();
exportProgressState.setBucket(bucket);
exportProgressState.setPrefix(prefix);
exportProgressState.setExportTime(exportTime.toString()); // information purpose
+ exportProgressState.setKmsKeyId(kmsKeyId);
ExportPartition exportPartition = new ExportPartition(tableArn, exportTime, Optional.of(exportProgressState));
coordinator.createPartition(exportPartition);
}
@@ -274,15 +294,13 @@ private TableInfo getTableInfo(TableConfig tableConfig) {
throw new InvalidPluginConfigurationException(errorMessage);
}
// Validate view type of DynamoDB stream
- if (describeTableResult.table().streamSpecification() != null) {
- String viewType = describeTableResult.table().streamSpecification().streamViewTypeAsString();
- LOG.debug("The stream view type for table " + tableName + " is " + viewType);
- List supportedType = List.of("NEW_IMAGE", "NEW_AND_OLD_IMAGES");
- if (!supportedType.contains(viewType)) {
- String errorMessage = "Stream " + tableConfig.getTableArn() + " is enabled with " + viewType + ". Supported types are " + supportedType;
- LOG.error(errorMessage);
- throw new InvalidPluginConfigurationException(errorMessage);
- }
+ String viewType = describeTableResult.table().streamSpecification().streamViewTypeAsString();
+ LOG.debug("The stream view type for table " + tableName + " is " + viewType);
+ List supportedType = List.of("NEW_IMAGE", "NEW_AND_OLD_IMAGES");
+ if (!supportedType.contains(viewType)) {
+ String errorMessage = "Stream " + tableConfig.getTableArn() + " is enabled with " + viewType + ". Supported types are " + supportedType;
+ LOG.error(errorMessage);
+ throw new InvalidPluginConfigurationException(errorMessage);
}
streamStartPosition = tableConfig.getStreamConfig().getStartPosition();
}
@@ -298,6 +316,7 @@ private TableInfo getTableInfo(TableConfig tableConfig) {
.streamStartPosition(streamStartPosition)
.exportBucket(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Bucket())
.exportPrefix(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Prefix())
+ .exportKmsKeyId(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3SseKmsKeyId())
.build();
return new TableInfo(tableConfig.getTableArn(), metadata);
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java
index ef35c82825..7ecd565ec4 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java
@@ -7,10 +7,10 @@
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
import org.opensearch.dataprepper.model.buffer.Buffer;
-import org.opensearch.dataprepper.model.configuration.PluginSetting;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.plugin.PluginFactory;
import org.opensearch.dataprepper.model.record.Record;
@@ -40,19 +40,26 @@ public class DynamoDBSource implements Source>, UsesEnhancedSource
private final ClientFactory clientFactory;
+ private final AcknowledgementSetManager acknowledgementSetManager;
+
private EnhancedSourceCoordinator coordinator;
private DynamoDBService dynamoDBService;
@DataPrepperPluginConstructor
- public DynamoDBSource(PluginMetrics pluginMetrics, final DynamoDBSourceConfig sourceConfig, final PluginFactory pluginFactory, final PluginSetting pluginSetting, final AwsCredentialsSupplier awsCredentialsSupplier) {
+ public DynamoDBSource(final PluginMetrics pluginMetrics,
+ final DynamoDBSourceConfig sourceConfig,
+ final PluginFactory pluginFactory,
+ final AwsCredentialsSupplier awsCredentialsSupplier,
+ final AcknowledgementSetManager acknowledgementSetManager) {
LOG.info("Create DynamoDB Source");
this.pluginMetrics = pluginMetrics;
this.sourceConfig = sourceConfig;
this.pluginFactory = pluginFactory;
+ this.acknowledgementSetManager = acknowledgementSetManager;
- clientFactory = new ClientFactory(awsCredentialsSupplier, sourceConfig.getAwsAuthenticationConfig());
+ clientFactory = new ClientFactory(awsCredentialsSupplier, sourceConfig.getAwsAuthenticationConfig(), sourceConfig.getTableConfigs().get(0).getExportConfig());
}
@Override
@@ -62,7 +69,7 @@ public void start(Buffer> buffer) {
coordinator.createPartition(new InitPartition());
// Create DynamoDB Service
- dynamoDBService = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics);
+ dynamoDBService = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics, acknowledgementSetManager);
dynamoDBService.init();
LOG.info("Start DynamoDB service");
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java
index 6e4a8fe7ae..2a682ce280 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java
@@ -7,10 +7,13 @@
import com.fasterxml.jackson.annotation.JsonProperty;
import jakarta.validation.Valid;
+import jakarta.validation.constraints.AssertTrue;
import jakarta.validation.constraints.NotNull;
import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.AwsAuthenticationConfig;
import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.TableConfig;
+import java.time.Duration;
+import java.util.Collections;
import java.util.List;
/**
@@ -19,14 +22,22 @@
public class DynamoDBSourceConfig {
@JsonProperty("tables")
- private List tableConfigs;
-
+ private List tableConfigs = Collections.emptyList();
@JsonProperty("aws")
@NotNull
@Valid
private AwsAuthenticationConfig awsAuthenticationConfig;
+ @JsonProperty("acknowledgments")
+ private boolean acknowledgments = false;
+
+ @JsonProperty("shard_acknowledgment_timeout")
+ private Duration shardAcknowledgmentTimeout = Duration.ofMinutes(3);
+
+ @JsonProperty("s3_data_file_acknowledgment_timeout")
+ private Duration dataFileAcknowledgmentTimeout = Duration.ofMinutes(5);
+
public DynamoDBSourceConfig() {
}
@@ -38,4 +49,19 @@ public List getTableConfigs() {
public AwsAuthenticationConfig getAwsAuthenticationConfig() {
return awsAuthenticationConfig;
}
+
+ public boolean isAcknowledgmentsEnabled() {
+ return acknowledgments;
+ }
+
+ public Duration getShardAcknowledgmentTimeout() {
+ return shardAcknowledgmentTimeout;
+ }
+
+ public Duration getDataFileAcknowledgmentTimeout() { return dataFileAcknowledgmentTimeout; }
+
+ @AssertTrue(message = "Exactly one table must be configured for the DynamoDb source.")
+ boolean isExactlyOneTableConfigured() {
+ return tableConfigs.size() == 1;
+ }
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java
index cb3463a3b6..dd1aac12a7 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java
@@ -6,7 +6,10 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.configuration;
import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.AssertTrue;
import jakarta.validation.constraints.NotBlank;
+import software.amazon.awssdk.arns.Arn;
+import software.amazon.awssdk.regions.Region;
public class ExportConfig {
@@ -16,6 +19,12 @@ public class ExportConfig {
@JsonProperty("s3_prefix")
private String s3Prefix;
+ @JsonProperty("s3_region")
+ private String s3Region;
+
+ @JsonProperty("s3_sse_kms_key_id")
+ private String s3SseKmsKeyId;
+
public String getS3Bucket() {
return s3Bucket;
}
@@ -23,4 +32,20 @@ public String getS3Bucket() {
public String getS3Prefix() {
return s3Prefix;
}
+
+ public Region getAwsRegion() {
+ return s3Region != null ? Region.of(s3Region) : null;
+ }
+
+ public String getS3SseKmsKeyId() {
+ return s3SseKmsKeyId;
+ }
+
+ @AssertTrue(message = "KMS Key ID must be a valid one.")
+ boolean isKmsKeyIdValid() {
+ // If key id is provided, it should be in a format like
+ // arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad3-80ca-63b12b3ec147
+ return s3SseKmsKeyId == null || Arn.fromString(s3SseKmsKeyId).resourceAsString() != null;
+ }
+
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java
index 70a6cbcf31..c259ccbcaa 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java
@@ -8,8 +8,9 @@
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.dataformat.ion.IonObjectMapper;
import io.micrometer.core.instrument.Counter;
+import org.opensearch.dataprepper.buffer.common.BufferAccumulator;
import org.opensearch.dataprepper.metrics.PluginMetrics;
-import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
@@ -18,7 +19,6 @@
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
public class ExportRecordConverter extends RecordConverter {
@@ -37,8 +37,8 @@ public class ExportRecordConverter extends RecordConverter {
private final Counter exportRecordSuccessCounter;
private final Counter exportRecordErrorCounter;
- public ExportRecordConverter(Buffer> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) {
- super(buffer, tableInfo);
+ public ExportRecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo, PluginMetrics pluginMetrics) {
+ super(bufferAccumulator, tableInfo);
this.pluginMetrics = pluginMetrics;
this.exportRecordSuccessCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSED_COUNT);
this.exportRecordErrorCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSING_ERROR_COUNT);
@@ -59,22 +59,27 @@ String getEventType() {
return "EXPORT";
}
- public void writeToBuffer(List lines) {
- List> data = lines.stream()
- .map(this::convertToMap)
- .map(d -> (Map) d.get(ITEM_KEY))
- .collect(Collectors.toList());
-
- List> events = data.stream().map(this::convertToEvent).collect(Collectors.toList());
+ public void writeToBuffer(final AcknowledgementSet acknowledgementSet, List lines) {
+
+ int eventCount = 0;
+ for (String line : lines) {
+ Map data = (Map) convertToMap(line).get(ITEM_KEY);
+ try {
+ addToBuffer(acknowledgementSet, data);
+ eventCount++;
+ } catch (Exception e) {
+ // will this cause too many logs?
+ LOG.error("Failed to add event to buffer due to {}", e.getMessage());
+ }
+ }
try {
- writeEventsToBuffer(events);
- exportRecordSuccessCounter.increment(events.size());
+ flushBuffer();
+ exportRecordSuccessCounter.increment(eventCount);
} catch (Exception e) {
- LOG.error("Failed to write {} events to buffer due to {}", events.size(), e.getMessage());
- exportRecordErrorCounter.increment(events.size());
+ LOG.error("Failed to write {} events to buffer due to {}", eventCount, e.getMessage());
+ exportRecordErrorCounter.increment(eventCount);
}
}
-
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java
index 0286627ba6..51c85bbefb 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java
@@ -6,15 +6,19 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.converter;
public class MetadataKeyAttributes {
- static final String PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "primary_key";
+ static final String PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "primary_key";
- static final String PARTITION_KEY_METADATA_ATTRIBUTE = "partition_key";
+ static final String PARTITION_KEY_METADATA_ATTRIBUTE = "partition_key";
- static final String SORT_KEY_METADATA_ATTRIBUTE = "sort_key";
+ static final String SORT_KEY_METADATA_ATTRIBUTE = "sort_key";
- static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "ts";
+ static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "dynamodb_timestamp";
- static final String STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "op";
+ static final String EVENT_DYNAMODB_ITEM_VERSION = "dynamodb_item_version";
- static final String EVENT_TABLE_NAME_METADATA_ATTRIBUTE = "table_name";
+ static final String EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "opensearch_action";
+
+ static final String DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE = "dynamodb_event_name";
+
+ static final String EVENT_TABLE_NAME_METADATA_ATTRIBUTE = "table_name";
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java
index cedf7fb0f1..7123b6825e 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java
@@ -5,7 +5,8 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.converter;
-import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.buffer.common.BufferAccumulator;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventMetadata;
import org.opensearch.dataprepper.model.event.JacksonEvent;
@@ -14,15 +15,16 @@
import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
import java.time.Instant;
-import java.util.List;
import java.util.Map;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_DYNAMODB_ITEM_VERSION;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE;
import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TABLE_NAME_METADATA_ATTRIBUTE;
import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE;
import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE;
import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.SORT_KEY_METADATA_ATTRIBUTE;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE;
/**
* Base Record Processor definition.
@@ -33,68 +35,91 @@ public abstract class RecordConverter {
private static final String DEFAULT_ACTION = OpenSearchBulkActions.INDEX.toString();
- private static final int DEFAULT_WRITE_TIMEOUT_MILLIS = 60_000;
- private final Buffer> buffer;
+ private final BufferAccumulator> bufferAccumulator;
private final TableInfo tableInfo;
- public RecordConverter(Buffer> buffer, TableInfo tableInfo) {
- this.buffer = buffer;
+ public RecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo) {
+ this.bufferAccumulator = bufferAccumulator;
this.tableInfo = tableInfo;
}
-
abstract String getEventType();
/**
- * Default method to conduct the document ID value,
- * Using partition key plus sort key (if any)
+ * Extract the value based on attribute map
+ *
+ * @param data A map of attribute name and value
+ * @param attributeName Attribute name
+ * @return the related attribute value, return null if the attribute name doesn't exist.
*/
- String getId(Map data) {
- String partitionKey = String.valueOf(data.get(tableInfo.getMetadata().getPartitionKeyAttributeName()));
- if (tableInfo.getMetadata().getSortKeyAttributeName() == null) {
- return partitionKey;
+ private String getAttributeValue(final Map data, String attributeName) {
+ if (data.containsKey(attributeName)) {
+ return String.valueOf(data.get(attributeName));
}
- String sortKey = String.valueOf(data.get(tableInfo.getMetadata().getSortKeyAttributeName()));
- return partitionKey + "_" + sortKey;
- }
-
- String getPartitionKey(final Map data) {
- return String.valueOf(data.get(tableInfo.getMetadata().getPartitionKeyAttributeName()));
- }
-
- String getSortKey(final Map data) {
- return String.valueOf(data.get(tableInfo.getMetadata().getSortKeyAttributeName()));
+ return null;
}
- void writeEventsToBuffer(List> events) throws Exception {
- buffer.writeAll(events, DEFAULT_WRITE_TIMEOUT_MILLIS);
+ void flushBuffer() throws Exception {
+ bufferAccumulator.flush();
}
- public Record convertToEvent(Map data, Instant eventCreationTime, String streamEventName) {
- Event event;
- event = JacksonEvent.builder()
+ /**
+ * Add event record to buffer
+ *
+ * @param data A map to hold event data, note that it may be empty.
+ * @param keys A map to hold the keys (partition key and sort key)
+ * @param eventCreationTimeMillis Creation timestamp of the event
+ * @param eventName Event name
+ * @throws Exception Exception if failed to write to buffer.
+ */
+ public void addToBuffer(final AcknowledgementSet acknowledgementSet,
+ final Map data,
+ final Map keys,
+ final long eventCreationTimeMillis,
+ final long eventVersionNumber,
+ final String eventName) throws Exception {
+ Event event = JacksonEvent.builder()
.withEventType(getEventType())
.withData(data)
.build();
+
+ // Only set external origination time for stream events, not export
+ if (eventName != null) {
+ final Instant externalOriginationTime = Instant.ofEpochMilli(eventCreationTimeMillis);
+ event.getEventHandle().setExternalOriginationTime(externalOriginationTime);
+ event.getMetadata().setExternalOriginationTime(externalOriginationTime);
+ }
EventMetadata eventMetadata = event.getMetadata();
eventMetadata.setAttribute(EVENT_TABLE_NAME_METADATA_ATTRIBUTE, tableInfo.getTableName());
- if (eventCreationTime != null) {
- eventMetadata.setAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE, eventCreationTime.toEpochMilli());
+ eventMetadata.setAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE, eventCreationTimeMillis);
+ eventMetadata.setAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE, eventName);
+ eventMetadata.setAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(eventName));
+ eventMetadata.setAttribute(EVENT_DYNAMODB_ITEM_VERSION, eventVersionNumber);
+
+ String partitionKey = getAttributeValue(keys, tableInfo.getMetadata().getPartitionKeyAttributeName());
+ eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, partitionKey);
+
+ String sortKey = getAttributeValue(keys, tableInfo.getMetadata().getSortKeyAttributeName());
+ if (sortKey != null) {
+ eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, sortKey);
+ eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey + "|" + sortKey);
+ } else {
+ eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey);
}
-
- eventMetadata.setAttribute(STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(streamEventName));
- eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, getId(data));
- eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, getPartitionKey(data));
- eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, getSortKey(data));
-
- return new Record<>(event);
+ if (acknowledgementSet != null) {
+ acknowledgementSet.add(event);
+ }
+ bufferAccumulator.add(new Record<>(event));
}
- public Record convertToEvent(Map data) {
- return convertToEvent(data, null, null);
+ public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map data) throws Exception {
+ // Export data doesn't have an event timestamp
+ // We consider this to be time of 0, meaning stream records will always be considered as newer
+ // than export records
+ addToBuffer(acknowledgementSet, data, data, System.currentTimeMillis(), 0L, null);
}
private String mapStreamEventNameToBulkAction(final String streamEventName) {
@@ -104,9 +129,8 @@ private String mapStreamEventNameToBulkAction(final String streamEventName) {
switch (streamEventName) {
case "INSERT":
- return OpenSearchBulkActions.CREATE.toString();
case "MODIFY":
- return OpenSearchBulkActions.UPSERT.toString();
+ return OpenSearchBulkActions.INDEX.toString();
case "REMOVE":
return OpenSearchBulkActions.DELETE.toString();
default:
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java
index bdfe6dcbe1..942e4d6a81 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java
@@ -9,18 +9,21 @@
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.micrometer.core.instrument.Counter;
+import org.opensearch.dataprepper.buffer.common.BufferAccumulator;
import org.opensearch.dataprepper.metrics.PluginMetrics;
-import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.event.Event;
-import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.Record;
+import java.time.Instant;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
public class StreamRecordConverter extends RecordConverter {
private static final Logger LOG = LoggerFactory.getLogger(StreamRecordConverter.class);
@@ -39,8 +42,11 @@ public class StreamRecordConverter extends RecordConverter {
private final Counter changeEventSuccessCounter;
private final Counter changeEventErrorCounter;
- public StreamRecordConverter(Buffer> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) {
- super(buffer, tableInfo);
+ private Instant currentSecond;
+ private int recordsSeenThisSecond = 0;
+
+ public StreamRecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo, PluginMetrics pluginMetrics) {
+ super(bufferAccumulator, tableInfo);
this.pluginMetrics = pluginMetrics;
this.changeEventSuccessCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSED_COUNT);
this.changeEventErrorCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSING_ERROR_COUNT);
@@ -51,30 +57,83 @@ String getEventType() {
return "STREAM";
}
- public void writeToBuffer(List records) {
- // TODO: What if convert failed.
- List> events = records.stream()
- .map(record -> convertToEvent(
- toMap(EnhancedDocument.fromAttributeValueMap(record.dynamodb().newImage()).toJson()),
- record.dynamodb().approximateCreationDateTime(),
- record.eventNameAsString()))
- .collect(Collectors.toList());
+
+ public void writeToBuffer(final AcknowledgementSet acknowledgementSet, List records) {
+
+ int eventCount = 0;
+ for (Record record : records) {
+ // NewImage may be empty
+ Map data = convertData(record.dynamodb().newImage());
+ // Always get keys from dynamodb().keys()
+ Map keys = convertKeys(record.dynamodb().keys());
+
+ try {
+ final long eventCreationTimeMillis = calculateTieBreakingVersionFromTimestamp(record.dynamodb().approximateCreationDateTime());
+ addToBuffer(acknowledgementSet, data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), eventCreationTimeMillis, record.eventNameAsString());
+ eventCount++;
+ } catch (Exception e) {
+ // will this cause too many logs?
+ LOG.error("Failed to add event to buffer due to {}", e.getMessage());
+ changeEventErrorCounter.increment();
+ }
+ }
try {
- writeEventsToBuffer(events);
- changeEventSuccessCounter.increment(events.size());
+ flushBuffer();
+ changeEventSuccessCounter.increment(eventCount);
} catch (Exception e) {
- LOG.error("Failed to write {} events to buffer due to {}", events.size(), e.getMessage());
- changeEventErrorCounter.increment(events.size());
+ LOG.error("Failed to write {} events to buffer due to {}", eventCount, e.getMessage());
+ changeEventErrorCounter.increment(eventCount);
}
}
- private Map toMap(String jsonData) {
+ /**
+ * Convert the DynamoDB attribute map to a normal map for data
+ */
+ private Map convertData(Map data) {
try {
+ String jsonData = EnhancedDocument.fromAttributeValueMap(data).toJson();
return MAPPER.readValue(jsonData, MAP_TYPE_REFERENCE);
} catch (JsonProcessingException e) {
return null;
}
}
+
+ /**
+ * Convert the DynamoDB attribute map to a normal map for keys
+ * This method may not be necessary, can use convertData() alternatively
+ */
+ private Map convertKeys(Map keys) {
+ Map result = new HashMap<>();
+ // The attribute type for key can only be N, B or S
+ keys.forEach(((attributeName, attributeValue) -> {
+ if (attributeValue.type() == AttributeValue.Type.N) {
+ // N for number
+ result.put(attributeName, attributeValue.n());
+ } else if (attributeValue.type() == AttributeValue.Type.B) {
+ // B for Binary
+ result.put(attributeName, attributeValue.b().toString());
+ } else {
+ result.put(attributeName, attributeValue.s());
+ }
+ }));
+ return result;
+
+ }
+
+ private long calculateTieBreakingVersionFromTimestamp(final Instant eventTimeInSeconds) {
+ if (currentSecond == null) {
+ currentSecond = eventTimeInSeconds;
+ } else if (currentSecond.isAfter(eventTimeInSeconds)) {
+ return eventTimeInSeconds.getEpochSecond() * 1_000_000;
+ } else if (currentSecond.isBefore(eventTimeInSeconds)) {
+ recordsSeenThisSecond = 0;
+ currentSecond = eventTimeInSeconds;
+ } else {
+ recordsSeenThisSecond++;
+ }
+
+ return eventTimeInSeconds.getEpochSecond() * 1_000_000 + recordsSeenThisSecond;
+ }
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java
index 63222b0d22..14422c6e8c 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java
@@ -23,7 +23,7 @@ public class DataFilePartition extends EnhancedSourcePartition> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT);
+ recordConverter = new ExportRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics);
+ this.acknowledgementSet = builder.acknowledgementSet;
+ this.dataFileAcknowledgmentTimeout = builder.dataFileAcknowledgmentTimeout;
}
- public static Builder builder() {
- return new Builder();
+ public static Builder builder(final S3ObjectReader s3ObjectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) {
+ return new Builder(s3ObjectReader, pluginMetrics, buffer);
}
@@ -69,9 +90,14 @@ public static Builder builder() {
*/
static class Builder {
- private S3ObjectReader s3ObjectReader;
+ private final S3ObjectReader objectReader;
+
+ private final PluginMetrics pluginMetrics;
+
+ private final Buffer> buffer;
+
+ private TableInfo tableInfo;
- private ExportRecordConverter recordConverter;
private DataFileCheckpointer checkpointer;
@@ -79,15 +105,20 @@ static class Builder {
private String key;
+ private AcknowledgementSet acknowledgementSet;
+
+ private Duration dataFileAcknowledgmentTimeout;
+
private int startLine;
- public Builder s3ObjectReader(S3ObjectReader s3ObjectReader) {
- this.s3ObjectReader = s3ObjectReader;
- return this;
+ public Builder(final S3ObjectReader objectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) {
+ this.objectReader = objectReader;
+ this.pluginMetrics = pluginMetrics;
+ this.buffer = buffer;
}
- public Builder recordConverter(ExportRecordConverter recordConverter) {
- this.recordConverter = recordConverter;
+ public Builder tableInfo(TableInfo tableInfo) {
+ this.tableInfo = tableInfo;
return this;
}
@@ -111,6 +142,16 @@ public Builder startLine(int startLine) {
return this;
}
+ public Builder acknowledgmentSet(AcknowledgementSet acknowledgementSet) {
+ this.acknowledgementSet = acknowledgementSet;
+ return this;
+ }
+
+ public Builder acknowledgmentSetTimeout(Duration dataFileAcknowledgmentTimeout) {
+ this.dataFileAcknowledgmentTimeout = dataFileAcknowledgmentTimeout;
+ return this;
+ }
+
public DataFileLoader build() {
return new DataFileLoader(this);
}
@@ -128,12 +169,12 @@ public void run() {
int lineCount = 0;
int lastLineProcessed = 0;
- try (GZIPInputStream gzipInputStream = new GZIPInputStream(s3ObjectReader.readFile(bucketName, key))) {
- BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream));
+ try (InputStream inputStream = objectReader.readFile(bucketName, key);
+ GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream);
+ BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream))) {
String line;
while ((line = reader.readLine()) != null) {
-
if (shouldStop) {
checkpointer.checkpoint(lastLineProcessed);
LOG.debug("Should Stop flag is set to True, looks like shutdown has triggered");
@@ -150,7 +191,7 @@ public void run() {
if ((lineCount - startLine) % DEFAULT_BATCH_SIZE == 0) {
// LOG.debug("Write to buffer for line " + (lineCount - DEFAULT_BATCH_SIZE) + " to " + lineCount);
- recordConverter.writeToBuffer(lines);
+ recordConverter.writeToBuffer(acknowledgementSet, lines);
lines.clear();
lastLineProcessed = lineCount;
}
@@ -165,16 +206,22 @@ public void run() {
}
if (!lines.isEmpty()) {
// Do final checkpoint.
- recordConverter.writeToBuffer(lines);
+ recordConverter.writeToBuffer(acknowledgementSet, lines);
checkpointer.checkpoint(lineCount);
}
lines.clear();
- reader.close();
- LOG.info("Complete loading s3://{}/{}", bucketName, key);
+
+ LOG.info("Completed loading s3://{}/{} to buffer", bucketName, key);
+
+ if (acknowledgementSet != null) {
+ checkpointer.updateDatafileForAcknowledgmentWait(dataFileAcknowledgmentTimeout);
+ acknowledgementSet.complete();
+ }
} catch (Exception e) {
checkpointer.checkpoint(lineCount);
- String errorMessage = String.format("Loading of s3://{}/{} completed with Exception: {}", bucketName, key, e.getMessage());
+
+ String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %s", bucketName, key, e.getMessage());
throw new RuntimeException(errorMessage);
}
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java
index 75dccf8218..7fee416bbc 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java
@@ -6,15 +6,17 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.export;
import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
-import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition;
import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
import software.amazon.awssdk.services.s3.S3Client;
+import java.time.Duration;
+
/**
* Factory class for DataFileLoader thread.
*/
@@ -22,32 +24,37 @@ public class DataFileLoaderFactory {
private final EnhancedSourceCoordinator coordinator;
- private final S3ObjectReader fileReader;
-
+ private final S3ObjectReader objectReader;
private final PluginMetrics pluginMetrics;
-
private final Buffer> buffer;
- public DataFileLoaderFactory(EnhancedSourceCoordinator coordinator, S3Client s3Client, PluginMetrics pluginMetrics, Buffer> buffer) {
+ public DataFileLoaderFactory(final EnhancedSourceCoordinator coordinator,
+ final S3Client s3Client,
+ final PluginMetrics pluginMetrics,
+ final Buffer> buffer) {
this.coordinator = coordinator;
this.pluginMetrics = pluginMetrics;
this.buffer = buffer;
- fileReader = new S3ObjectReader(s3Client);
+ objectReader = new S3ObjectReader(s3Client);
}
- public Runnable createDataFileLoader(DataFilePartition dataFilePartition, TableInfo tableInfo) {
- ExportRecordConverter recordProcessor = new ExportRecordConverter(buffer, tableInfo, pluginMetrics);
+ public Runnable createDataFileLoader(final DataFilePartition dataFilePartition,
+ final TableInfo tableInfo,
+ final AcknowledgementSet acknowledgementSet,
+ final Duration acknowledgmentTimeout) {
DataFileCheckpointer checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition);
// Start a data loader thread.
- DataFileLoader loader = DataFileLoader.builder()
- .s3ObjectReader(fileReader)
+ DataFileLoader loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer)
.bucketName(dataFilePartition.getBucket())
.key(dataFilePartition.getKey())
- .recordConverter(recordProcessor)
+ .tableInfo(tableInfo)
.checkpointer(checkpointer)
- .startLine(dataFilePartition.getProgressState().get().getLoaded())
+ .acknowledgmentSet(acknowledgementSet)
+ .acknowledgmentSetTimeout(acknowledgmentTimeout)
+ // We can't checkpoint with acks enabled yet
+ .startLine(acknowledgementSet == null ? dataFilePartition.getProgressState().get().getLoaded() : 0)
.build();
return loader;
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java
index 6c56323cca..f02cf2f357 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java
@@ -7,8 +7,11 @@
import io.micrometer.core.instrument.Counter;
import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition;
+import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState;
import org.opensearch.dataprepper.plugins.source.dynamodb.model.LoadStatus;
@@ -35,7 +38,7 @@ public class DataFileScheduler implements Runnable {
/**
* Maximum concurrent data loader per node
*/
- private static final int MAX_JOB_COUNT = 4;
+ private static final int MAX_JOB_COUNT = 3;
/**
* Default interval to acquire a lease from coordination store
@@ -54,16 +57,25 @@ public class DataFileScheduler implements Runnable {
private final PluginMetrics pluginMetrics;
+ private final AcknowledgementSetManager acknowledgementSetManager;
+
+ private final DynamoDBSourceConfig dynamoDBSourceConfig;
+
private final Counter exportFileSuccessCounter;
private final AtomicLong activeExportS3ObjectConsumersGauge;
- public DataFileScheduler(EnhancedSourceCoordinator coordinator, DataFileLoaderFactory loaderFactory, PluginMetrics pluginMetrics) {
+ public DataFileScheduler(final EnhancedSourceCoordinator coordinator,
+ final DataFileLoaderFactory loaderFactory,
+ final PluginMetrics pluginMetrics,
+ final AcknowledgementSetManager acknowledgementSetManager,
+ final DynamoDBSourceConfig dynamoDBSourceConfig) {
this.coordinator = coordinator;
this.pluginMetrics = pluginMetrics;
this.loaderFactory = loaderFactory;
-
+ this.acknowledgementSetManager = acknowledgementSetManager;
+ this.dynamoDBSourceConfig = dynamoDBSourceConfig;
executor = Executors.newFixedThreadPool(MAX_JOB_COUNT);
@@ -77,34 +89,63 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) {
TableInfo tableInfo = getTableInfo(tableArn);
- Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo);
+ final boolean acknowledgmentsEnabled = dynamoDBSourceConfig.isAcknowledgmentsEnabled();
+
+ AcknowledgementSet acknowledgementSet = null;
+ if (acknowledgmentsEnabled) {
+ acknowledgementSet = acknowledgementSetManager.create((result) -> {
+ if (result == true) {
+ completeDataLoader(dataFilePartition).accept(null, null);
+ LOG.info("Received acknowledgment of completion from sink for data file {}", dataFilePartition.getKey());
+ } else {
+ LOG.warn("Negative acknowledgment received for data file {}, retrying", dataFilePartition.getKey());
+ coordinator.giveUpPartition(dataFilePartition);
+ }
+ }, dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout());
+ }
+
+ Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo, acknowledgementSet, dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout());
CompletableFuture runLoader = CompletableFuture.runAsync(loader, executor);
- runLoader.whenComplete(completeDataLoader(dataFilePartition));
+
+ if (!acknowledgmentsEnabled) {
+ runLoader.whenComplete(completeDataLoader(dataFilePartition));
+ } else {
+ runLoader.whenComplete((v, ex) -> numOfWorkers.decrementAndGet());
+ }
numOfWorkers.incrementAndGet();
}
@Override
public void run() {
- LOG.info("Start running Data File Scheduler");
+ LOG.debug("Starting Data File Scheduler to process S3 data files for export");
while (!Thread.currentThread().isInterrupted()) {
- if (numOfWorkers.get() < MAX_JOB_COUNT) {
- final Optional sourcePartition = coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE);
-
- if (sourcePartition.isPresent()) {
- activeExportS3ObjectConsumersGauge.incrementAndGet();
- DataFilePartition dataFilePartition = (DataFilePartition) sourcePartition.get();
- processDataFilePartition(dataFilePartition);
- activeExportS3ObjectConsumersGauge.decrementAndGet();
- }
- }
try {
- Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS);
- } catch (final InterruptedException e) {
- LOG.info("InterruptedException occurred");
- break;
+ if (numOfWorkers.get() < MAX_JOB_COUNT) {
+ final Optional sourcePartition = coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE);
+
+ if (sourcePartition.isPresent()) {
+ activeExportS3ObjectConsumersGauge.incrementAndGet();
+ DataFilePartition dataFilePartition = (DataFilePartition) sourcePartition.get();
+ processDataFilePartition(dataFilePartition);
+ activeExportS3ObjectConsumersGauge.decrementAndGet();
+ }
+ }
+ try {
+ Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS);
+ } catch (final InterruptedException e) {
+ LOG.info("The DataFileScheduler was interrupted while waiting to retry, stopping processing");
+ break;
+ }
+ } catch (final Exception e) {
+ LOG.error("Received an exception while processing an S3 data file, backing off and retrying", e);
+ try {
+ Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS);
+ } catch (final InterruptedException ex) {
+ LOG.info("The DataFileScheduler was interrupted while waiting to retry, stopping processing");
+ break;
+ }
}
-
}
LOG.warn("Data file scheduler is interrupted, Stop all data file loaders...");
// Cannot call executor.shutdownNow() here
@@ -138,7 +179,10 @@ private String getStreamArn(String exportArn) {
private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) {
return (v, ex) -> {
- numOfWorkers.decrementAndGet();
+
+ if (!dynamoDBSourceConfig.isAcknowledgmentsEnabled()) {
+ numOfWorkers.decrementAndGet();
+ }
if (ex == null) {
exportFileSuccessCounter.increment();
// Update global state
@@ -148,8 +192,7 @@ private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) {
} else {
// The data loader must have already done one last checkpointing.
- LOG.debug("Data Loader completed with exception");
- LOG.error("{}", ex);
+ LOG.error("Loading S3 data files completed with an exception: {}", ex);
// Release the ownership
coordinator.giveUpPartition(dataFilePartition);
}
@@ -157,6 +200,17 @@ private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) {
};
}
+ /**
+ * There is a global state with sourcePartitionKey the export Arn,
+ * to track the number of files are processed.
+ * Each time, load of a data file is completed,
+ * The state must be updated.
+ * Note that the state may be updated since multiple threads are updating the same state.
+ * Retry is required.
+ *
+ * @param exportArn Export Arn.
+ * @param loaded Number records Loaded.
+ */
private void updateState(String exportArn, int loaded) {
String streamArn = getStreamArn(exportArn);
@@ -172,27 +226,24 @@ private void updateState(String exportArn, int loaded) {
GlobalState globalState = (GlobalState) globalPartition.get();
LoadStatus loadStatus = LoadStatus.fromMap(globalState.getProgressState().get());
- LOG.debug("Current status: total {} loaded {}", loadStatus.getTotalFiles(), loadStatus.getLoadedFiles());
-
loadStatus.setLoadedFiles(loadStatus.getLoadedFiles() + 1);
+ LOG.info("Current status: total {} loaded {}", loadStatus.getTotalFiles(), loadStatus.getLoadedFiles());
+
loadStatus.setLoadedRecords(loadStatus.getLoadedRecords() + loaded);
globalState.setProgressState(loadStatus.toMap());
try {
- coordinator.saveProgressStateForPartition(globalState);
+ coordinator.saveProgressStateForPartition(globalState, null);
// if all load are completed.
if (streamArn != null && loadStatus.getLoadedFiles() == loadStatus.getTotalFiles()) {
- LOG.debug("All Exports are done, streaming can continue...");
+ LOG.info("All Exports are done, streaming can continue...");
coordinator.createPartition(new GlobalState(streamArn, Optional.empty()));
}
break;
} catch (Exception e) {
- LOG.error("Failed to update the global status, looks like the status was out of dated, will retry..");
+ LOG.error("Failed to update the global status, looks like the status was out of date, will retry..");
}
-
}
-
-
}
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java
index 93f9e5b51a..c150643d56 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java
@@ -94,32 +94,41 @@ public ExportScheduler(EnhancedSourceCoordinator enhancedSourceCoordinator, Dyna
public void run() {
LOG.debug("Start running Export Scheduler");
while (!Thread.currentThread().isInterrupted()) {
- // Does not have limit on max leases
- // As most of the time it's just to wait
- final Optional sourcePartition = enhancedSourceCoordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE);
+ try {
+ // Does not have limit on max leases
+ // As most of the time it's just to wait
+ final Optional sourcePartition = enhancedSourceCoordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE);
- if (sourcePartition.isPresent()) {
+ if (sourcePartition.isPresent()) {
- ExportPartition exportPartition = (ExportPartition) sourcePartition.get();
- LOG.debug("Acquired an export partition: " + exportPartition.getPartitionKey());
+ ExportPartition exportPartition = (ExportPartition) sourcePartition.get();
+ LOG.debug("Acquired an export partition: " + exportPartition.getPartitionKey());
- String exportArn = getOrCreateExportArn(exportPartition);
+ String exportArn = getOrCreateExportArn(exportPartition);
- if (exportArn == null) {
- closeExportPartitionWithError(exportPartition);
- } else {
- CompletableFuture checkStatus = CompletableFuture.supplyAsync(() -> checkExportStatus(exportPartition), executor);
- checkStatus.whenComplete(completeExport(exportPartition));
- }
+ if (exportArn == null) {
+ closeExportPartitionWithError(exportPartition);
+ } else {
+ CompletableFuture checkStatus = CompletableFuture.supplyAsync(() -> checkExportStatus(exportPartition), executor);
+ checkStatus.whenComplete(completeExport(exportPartition));
+ }
+ }
+ try {
+ Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS);
+ } catch (final InterruptedException e) {
+ LOG.info("The ExportScheduler was interrupted while waiting to retry, stopping processing");
+ break;
+ }
+ } catch (final Exception e) {
+ LOG.error("Received an exception during export from DynamoDB to S3, backing off and retrying", e);
+ try {
+ Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS);
+ } catch (final InterruptedException ex) {
+ LOG.info("The ExportScheduler was interrupted while waiting to retry, stopping processing");
+ break;
+ }
}
- try {
- Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS);
- } catch (final InterruptedException e) {
- LOG.info("InterruptedException occurred");
- break;
- }
-
}
LOG.warn("Export scheduler interrupted, looks like shutdown has triggered");
executor.shutdownNow();
@@ -130,7 +139,7 @@ public void run() {
private BiConsumer completeExport(ExportPartition exportPartition) {
return (status, ex) -> {
if (ex != null) {
- LOG.debug("Check export status for {} failed with error {}", exportPartition.getPartitionKey(), ex.getMessage());
+ LOG.warn("Check export status for {} failed with error {}", exportPartition.getPartitionKey(), ex.getMessage());
// closeExportPartitionWithError(exportPartition);
enhancedSourceCoordinator.giveUpPartition(exportPartition);
} else {
@@ -147,8 +156,7 @@ private BiConsumer completeExport(ExportPartition exportParti
ExportProgressState state = exportPartition.getProgressState().get();
String bucketName = state.getBucket();
String exportArn = state.getExportArn();
-
-
+
String manifestKey = exportTaskManager.getExportManifest(exportArn);
LOG.debug("Export manifest summary file is " + manifestKey);
@@ -173,13 +181,14 @@ private BiConsumer completeExport(ExportPartition exportParti
private void createDataFilePartitions(String exportArn, String bucketName, Map dataFileInfo) {
- LOG.info("Totally {} data files generated for export {}", dataFileInfo.size(), exportArn);
+ LOG.info("Total of {} data files generated for export {}", dataFileInfo.size(), exportArn);
AtomicInteger totalRecords = new AtomicInteger();
AtomicInteger totalFiles = new AtomicInteger();
dataFileInfo.forEach((key, size) -> {
DataFileProgressState progressState = new DataFileProgressState();
progressState.setTotal(size);
progressState.setLoaded(0);
+
totalFiles.addAndGet(1);
totalRecords.addAndGet(size);
DataFilePartition partition = new DataFilePartition(exportArn, bucketName, key, Optional.of(progressState));
@@ -197,6 +206,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map DEFAULT_CHECKPOINT_INTERVAL_MILLS) {
- enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition);
+ enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition, null);
lastCheckpointTime = System.currentTimeMillis();
}
@@ -250,14 +260,14 @@ private String getOrCreateExportArn(ExportPartition exportPartition) {
LOG.info("Try to submit a new export job for table {} with export time {}", exportPartition.getTableArn(), exportPartition.getExportTime());
// submit a new export request
- String exportArn = exportTaskManager.submitExportJob(exportPartition.getTableArn(), state.getBucket(), state.getPrefix(), exportPartition.getExportTime());
+ String exportArn = exportTaskManager.submitExportJob(exportPartition.getTableArn(), state.getBucket(), state.getPrefix(), state.getKmsKeyId(), exportPartition.getExportTime());
// Update state with export Arn in the coordination table.
// So that it won't be submitted again after a restart.
if (exportArn != null) {
LOG.info("Export arn is " + exportArn);
state.setExportArn(exportArn);
- enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition);
+ enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition, null);
}
return exportArn;
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportTaskManager.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportTaskManager.java
index c0b43009a0..95d6bce986 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportTaskManager.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportTaskManager.java
@@ -14,6 +14,7 @@
import software.amazon.awssdk.services.dynamodb.model.ExportFormat;
import software.amazon.awssdk.services.dynamodb.model.ExportTableToPointInTimeRequest;
import software.amazon.awssdk.services.dynamodb.model.ExportTableToPointInTimeResponse;
+import software.amazon.awssdk.services.dynamodb.model.S3SseAlgorithm;
import java.time.Instant;
@@ -30,12 +31,15 @@ public ExportTaskManager(DynamoDbClient dynamoDBClient) {
this.dynamoDBClient = dynamoDBClient;
}
- public String submitExportJob(String tableArn, String bucketName, String prefix, Instant exportTime) {
+ public String submitExportJob(String tableArn, String bucket, String prefix, String kmsKeyId, Instant exportTime) {
+ S3SseAlgorithm algorithm = kmsKeyId == null || kmsKeyId.isEmpty() ? S3SseAlgorithm.AES256 : S3SseAlgorithm.KMS;
// No needs to use a client token here.
ExportTableToPointInTimeRequest req = ExportTableToPointInTimeRequest.builder()
.tableArn(tableArn)
- .s3Bucket(bucketName)
+ .s3Bucket(bucket)
.s3Prefix(prefix)
+ .s3SseAlgorithm(algorithm)
+ .s3SseKmsKeyId(kmsKeyId)
.exportFormat(DEFAULT_EXPORT_FORMAT)
.exportTime(exportTime)
.build();
@@ -46,7 +50,6 @@ public String submitExportJob(String tableArn, String bucketName, String prefix,
String exportArn = response.exportDescription().exportArn();
String status = response.exportDescription().exportStatusAsString();
-
LOG.debug("Export Job submitted with ARN {} and status {}", exportArn, status);
return exportArn;
} catch (SdkException e) {
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReader.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReader.java
index 00919f1a5d..6f1d16847a 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReader.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReader.java
@@ -5,7 +5,6 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.export;
-import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.opensearch.dataprepper.plugins.source.dynamodb.model.ExportSummary;
@@ -13,7 +12,6 @@
import org.slf4j.LoggerFactory;
import java.io.BufferedReader;
-import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashMap;
@@ -39,22 +37,18 @@ public ManifestFileReader(S3ObjectReader objectReader) {
public ExportSummary parseSummaryFile(String bucket, String key) {
LOG.debug("Try to read the manifest summary file");
- InputStream object = objectReader.readFile(bucket, key);
- BufferedReader reader = new BufferedReader(new InputStreamReader(object));
- try {
+ try (InputStream object = objectReader.readFile(bucket, key);
+ BufferedReader reader = new BufferedReader(new InputStreamReader(object))) {
+ // Only one line
String line = reader.readLine();
LOG.debug("Manifest summary: {}", line);
ExportSummary summaryInfo = MAPPER.readValue(line, ExportSummary.class);
return summaryInfo;
- } catch (JsonProcessingException e) {
+ } catch (Exception e) {
LOG.error("Failed to parse the summary info due to {}", e.getMessage());
throw new RuntimeException(e);
-
- } catch (IOException e) {
- LOG.error("IO Exception due to {}", e.getMessage());
- throw new RuntimeException(e);
}
}
@@ -63,11 +57,10 @@ public Map parseDataFile(String bucket, String key) {
LOG.info("Try to read the manifest data file");
Map result = new HashMap<>();
- InputStream object = objectReader.readFile(bucket, key);
- BufferedReader reader = new BufferedReader(new InputStreamReader(object));
-
- String line;
- try {
+
+ try (InputStream object = objectReader.readFile(bucket, key);
+ BufferedReader reader = new BufferedReader(new InputStreamReader(object))) {
+ String line;
while ((line = reader.readLine()) != null) {
// An example line as below:
// {"itemCount":46331,"md5Checksum":"a0k21IY3eelgr2PuWJLjJw==","etag":"51f9f394903c5d682321c6211aae8b6a-1","dataFileS3Key":"test-table-export/AWSDynamoDB/01692350182719-6de2c037/data/fpgzwz7ome3s7a5gqn2mu3ogtq.json.gz"}
@@ -76,8 +69,9 @@ public Map parseDataFile(String bucket, String key) {
result.put(map.get(DATA_FILE_S3_KEY), Integer.valueOf(map.get(DATA_FILE_ITEM_COUNT_KEY)));
}
- } catch (IOException e) {
- LOG.error("IO Exception due to {}", e.getMessage());
+ } catch (Exception e) {
+ LOG.error("Exception due to {}", e.getMessage());
+ throw new RuntimeException(e);
}
return result;
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java
index 881f59f605..578c5c7625 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java
@@ -57,7 +57,7 @@ public class ExportSummary {
private long billedSizeBytes;
@JsonProperty("itemCount")
- private int itemCount;
+ private long itemCount;
@JsonProperty("outputFormat")
private String outputFormat;
@@ -115,7 +115,7 @@ public long getBilledSizeBytes() {
return billedSizeBytes;
}
- public int getItemCount() {
+ public long getItemCount() {
return itemCount;
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java
index fd84c87e98..6493c70b51 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java
@@ -18,11 +18,11 @@ public class LoadStatus {
private int loadedFiles;
- private int totalRecords;
+ private long totalRecords;
- private int loadedRecords;
+ private long loadedRecords;
- public LoadStatus(int totalFiles, int loadedFiles, int totalRecords, int loadedRecords) {
+ public LoadStatus(int totalFiles, int loadedFiles, long totalRecords, long loadedRecords) {
this.totalFiles = totalFiles;
this.loadedFiles = loadedFiles;
this.totalRecords = totalRecords;
@@ -45,7 +45,7 @@ public void setLoadedFiles(int loadedFiles) {
this.loadedFiles = loadedFiles;
}
- public int getTotalRecords() {
+ public long getTotalRecords() {
return totalRecords;
}
@@ -53,11 +53,11 @@ public void setTotalRecords(int totalRecords) {
this.totalRecords = totalRecords;
}
- public int getLoadedRecords() {
+ public long getLoadedRecords() {
return loadedRecords;
}
- public void setLoadedRecords(int loadedRecords) {
+ public void setLoadedRecords(long loadedRecords) {
this.loadedRecords = loadedRecords;
}
@@ -72,10 +72,10 @@ public Map toMap() {
public static LoadStatus fromMap(Map map) {
return new LoadStatus(
- (int) map.get(TOTAL_FILES),
- (int) map.get(LOADED_FILES),
- (int) map.get(TOTAL_RECORDS),
- (int) map.get(LOADED_RECORDS)
+ ((Number) map.get(TOTAL_FILES)).intValue(),
+ ((Number) map.get(LOADED_FILES)).intValue(),
+ ((Number) map.get(TOTAL_RECORDS)).longValue(),
+ ((Number) map.get(LOADED_RECORDS)).longValue()
);
}
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java
index 8415a41001..ba21304d31 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java
@@ -19,7 +19,6 @@ public class TableMetadata {
private static final String REQUIRE_EXPORT_KEY = "export";
private static final String REQUIRE_STREAM_KEY = "stream";
-
private final String partitionKeyAttributeName;
private final String sortKeyAttributeName;
@@ -36,6 +35,8 @@ public class TableMetadata {
private final String exportPrefix;
+ private final String exportKmsKeyId;
+
private TableMetadata(Builder builder) {
this.partitionKeyAttributeName = builder.partitionKeyAttributeName;
this.sortKeyAttributeName = builder.sortKeyAttributeName;
@@ -45,6 +46,7 @@ private TableMetadata(Builder builder) {
this.exportBucket = builder.exportBucket;
this.exportPrefix = builder.exportPrefix;
this.streamStartPosition = builder.streamStartPosition;
+ this.exportKmsKeyId = builder.exportKmsKeyId;
}
@@ -70,6 +72,8 @@ public static class Builder {
private String exportPrefix;
+ private String exportKmsKeyId;
+
private StreamStartPosition streamStartPosition;
@@ -108,6 +112,11 @@ public Builder exportPrefix(String exportPrefix) {
return this;
}
+ public Builder exportKmsKeyId(String exportKmsKeyId) {
+ this.exportKmsKeyId = exportKmsKeyId;
+ return this;
+ }
+
public Builder streamStartPosition(StreamStartPosition streamStartPosition) {
this.streamStartPosition = streamStartPosition;
return this;
@@ -173,4 +182,8 @@ public String getExportBucket() {
public String getExportPrefix() {
return exportPrefix;
}
+
+ public String getExportKmsKeyId() {
+ return exportKmsKeyId;
+ }
}
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java
index ef47eb39ef..c5da709775 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java
@@ -5,16 +5,22 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
-
+import org.opensearch.dataprepper.buffer.common.BufferAccumulator;
+import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import org.opensearch.dataprepper.model.buffer.Buffer;
+import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter;
+import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.core.exception.SdkException;
import software.amazon.awssdk.services.dynamodb.model.GetRecordsRequest;
import software.amazon.awssdk.services.dynamodb.model.GetRecordsResponse;
-import software.amazon.awssdk.services.dynamodb.model.Record;
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient;
+import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.stream.Collectors;
@@ -41,16 +47,36 @@ public class ShardConsumer implements Runnable {
*/
private static final int GET_RECORD_INTERVAL_MILLS = 300;
+ /**
+ * Idle Time between GetRecords Reads
+ */
+ private static final int MINIMUM_GET_RECORD_INTERVAL_MILLS = 10;
+
+ /**
+ * Minimum Idle Time between GetRecords Reads
+ */
+ private static final long GET_RECORD_DELAY_THRESHOLD_MILLS = 15_000;
+
/**
* Default interval to check if export is completed.
*/
private static final int DEFAULT_WAIT_FOR_EXPORT_INTERVAL_MILLS = 60_000;
+
+ /**
+ * Default number of times in the wait for export to do regular checkpoint.
+ */
+ private static final int DEFAULT_WAIT_COUNT_TO_CHECKPOINT = 5;
+
/**
* Default regular checkpoint interval
*/
private static final int DEFAULT_CHECKPOINT_INTERVAL_MILLS = 2 * 60_000;
+ static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60);
+ static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000;
+
+
private final DynamoDbStreamsClient dynamoDbStreamsClient;
private final StreamRecordConverter recordConverter;
@@ -63,17 +89,24 @@ public class ShardConsumer implements Runnable {
private boolean waitForExport;
+ private final AcknowledgementSet acknowledgementSet;
+
+ private final Duration shardAcknowledgmentTimeout;
+
private ShardConsumer(Builder builder) {
this.dynamoDbStreamsClient = builder.dynamoDbStreamsClient;
- this.recordConverter = builder.recordConverter;
this.checkpointer = builder.checkpointer;
this.shardIterator = builder.shardIterator;
this.startTime = builder.startTime;
this.waitForExport = builder.waitForExport;
+ final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT);
+ recordConverter = new StreamRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics);
+ this.acknowledgementSet = builder.acknowledgementSet;
+ this.shardAcknowledgmentTimeout = builder.dataFileAcknowledgmentTimeout;
}
- public static Builder builder(DynamoDbStreamsClient streamsClient) {
- return new Builder(streamsClient);
+ public static Builder builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) {
+ return new Builder(dynamoDbStreamsClient, pluginMetrics, buffer);
}
@@ -81,8 +114,11 @@ static class Builder {
private final DynamoDbStreamsClient dynamoDbStreamsClient;
+ private final PluginMetrics pluginMetrics;
- private StreamRecordConverter recordConverter;
+ private final Buffer> buffer;
+
+ private TableInfo tableInfo;
private StreamCheckpointer checkpointer;
@@ -93,12 +129,17 @@ static class Builder {
private boolean waitForExport;
- public Builder(DynamoDbStreamsClient dynamoDbStreamsClient) {
+ private AcknowledgementSet acknowledgementSet;
+ private Duration dataFileAcknowledgmentTimeout;
+
+ public Builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) {
this.dynamoDbStreamsClient = dynamoDbStreamsClient;
+ this.pluginMetrics = pluginMetrics;
+ this.buffer = buffer;
}
- public Builder recordConverter(StreamRecordConverter recordConverter) {
- this.recordConverter = recordConverter;
+ public Builder tableInfo(TableInfo tableInfo) {
+ this.tableInfo = tableInfo;
return this;
}
@@ -122,6 +163,16 @@ public Builder waitForExport(boolean waitForExport) {
return this;
}
+ public Builder acknowledgmentSet(AcknowledgementSet acknowledgementSet) {
+ this.acknowledgementSet = acknowledgementSet;
+ return this;
+ }
+
+ public Builder acknowledgmentSetTimeout(Duration dataFileAcknowledgmentTimeout) {
+ this.dataFileAcknowledgmentTimeout = dataFileAcknowledgmentTimeout;
+ return this;
+ }
+
public ShardConsumer build() {
return new ShardConsumer(this);
}
@@ -131,7 +182,7 @@ public ShardConsumer build() {
@Override
public void run() {
- LOG.info("Shard Consumer start to run...");
+ LOG.debug("Shard Consumer start to run...");
long lastCheckpointTime = System.currentTimeMillis();
String sequenceNumber = "";
@@ -139,7 +190,7 @@ public void run() {
while (!shouldStop) {
if (shardIterator == null) {
// End of Shard
- LOG.debug("Reach end of shard");
+ LOG.debug("Reached end of shard");
checkpointer.checkpoint(sequenceNumber);
break;
}
@@ -157,7 +208,7 @@ public void run() {
.build();
- List records;
+ List records;
GetRecordsResponse response;
try {
response = dynamoDbStreamsClient.getRecords(req);
@@ -168,12 +219,15 @@ public void run() {
shardIterator = response.nextShardIterator();
+ int interval;
+
if (!response.records().isEmpty()) {
// Always use the last sequence number for checkpoint
sequenceNumber = response.records().get(response.records().size() - 1).dynamodb().sequenceNumber();
+ Instant lastEventTime = response.records().get(response.records().size() - 1).dynamodb().approximateCreationDateTime();
if (waitForExport) {
- Instant lastEventTime = response.records().get(response.records().size() - 1).dynamodb().approximateCreationDateTime();
+
if (lastEventTime.compareTo(startTime) <= 0) {
LOG.debug("Get {} events before start time, ignore...", response.records().size());
continue;
@@ -188,16 +242,26 @@ public void run() {
} else {
records = response.records();
}
- recordConverter.writeToBuffer(records);
+ recordConverter.writeToBuffer(acknowledgementSet, records);
+ long delay = System.currentTimeMillis() - lastEventTime.toEpochMilli();
+ interval = delay > GET_RECORD_DELAY_THRESHOLD_MILLS ? MINIMUM_GET_RECORD_INTERVAL_MILLS : GET_RECORD_INTERVAL_MILLS;
+
+ } else {
+ interval = GET_RECORD_INTERVAL_MILLS;
}
try {
// Idle between get records call.
- Thread.sleep(GET_RECORD_INTERVAL_MILLS);
+ Thread.sleep(interval);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
}
+ if (acknowledgementSet != null) {
+ checkpointer.updateShardForAcknowledgmentWait(shardAcknowledgmentTimeout);
+ acknowledgementSet.complete();
+ }
+
// interrupted
if (shouldStop) {
// Do last checkpoint and then quit
@@ -205,14 +269,27 @@ public void run() {
checkpointer.checkpoint(sequenceNumber);
throw new RuntimeException("Shard Consumer is interrupted");
}
+
+ if (waitForExport) {
+ waitForExport();
+ }
}
private void waitForExport() {
LOG.debug("Start waiting for export to be done and loaded");
+ int numberOfWaits = 0;
while (!checkpointer.isExportDone()) {
LOG.debug("Export is in progress, wait...");
try {
Thread.sleep(DEFAULT_WAIT_FOR_EXPORT_INTERVAL_MILLS);
+ // The wait for export may take a long time
+ // Need to extend the timeout of the ownership in the coordination store.
+ // Otherwise, the lease will expire.
+ numberOfWaits++;
+ if (numberOfWaits % DEFAULT_WAIT_COUNT_TO_CHECKPOINT == 0) {
+ // To extend the timeout of lease
+ checkpointer.checkpoint(null);
+ }
} catch (InterruptedException e) {
LOG.error("Wait for export is interrupted ({})", e.getMessage());
// Directly quit the process
diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java
index 1c35f882ad..7cdceba3ff 100644
--- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java
+++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java
@@ -6,11 +6,11 @@
package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
import org.opensearch.dataprepper.metrics.PluginMetrics;
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
-import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition;
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
@@ -20,6 +20,7 @@
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient;
+import java.time.Duration;
import java.time.Instant;
import java.util.Optional;
@@ -31,17 +32,19 @@ public class ShardConsumerFactory {
private static final int STREAM_TO_TABLE_OFFSET = "stream/".length();
+
private final DynamoDbStreamsClient streamsClient;
private final EnhancedSourceCoordinator enhancedSourceCoordinator;
private final PluginMetrics pluginMetrics;
private final ShardManager shardManager;
-
private final Buffer> buffer;
+
public ShardConsumerFactory(final EnhancedSourceCoordinator enhancedSourceCoordinator,
final DynamoDbStreamsClient streamsClient, final PluginMetrics pluginMetrics,
- final ShardManager shardManager, final Buffer> buffer) {
+ final ShardManager shardManager,
+ final Buffer