Skip to content
New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Add aggregate_threshold with maximum_size to s3 sink #4385

Merged
merged 3 commits into from
Apr 4, 2024
Merged

Add aggregate_threshold with maximum_size to s3 sink #4385

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7d052d9
Add aggregate_threshold with maximum_size to s3 sink
graytaylor0 Apr 1, 2024
4709b66
Address PR comments
graytaylor0 Apr 4, 2024
ec98dc0
Address PR comment
graytaylor0 Apr 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions ...integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/NdjsonOutputScenario.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,11 +68,12 @@ public void validateDynamicPartition(int expectedRecords, int partitionNumber, f

final Map<String, Object> actualData = OBJECT_MAPPER.readValue(actualJsonString, Map.class);
assertThat(actualData.get("sequence"), equalTo(partitionNumber));

count++;
}

assertThat(count, equalTo(expectedRecords));
if (expectedRecords != -1) {
assertThat(count, equalTo(expectedRecords));
}
}

@Override
Expand Down
98 changes: 97 additions & 1 deletion ...s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
import org.opensearch.dataprepper.expression.ExpressionEvaluator;
import org.opensearch.dataprepper.model.codec.OutputCodec;
import org.opensearch.dataprepper.model.configuration.PluginModel;
import org.opensearch.dataprepper.model.configuration.PluginSetting;
Expand All @@ -32,6 +33,7 @@
import org.opensearch.dataprepper.model.sink.SinkContext;
import org.opensearch.dataprepper.model.types.ByteCount;
import org.opensearch.dataprepper.plugins.sink.s3.accumulator.BufferTypeOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.AggregateThresholdOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.AwsAuthenticationOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.ObjectKeyOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.ThresholdOptions;
Expand All @@ -45,7 +47,6 @@
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
import software.amazon.awssdk.services.s3.model.S3Object;
import org.opensearch.dataprepper.expression.ExpressionEvaluator;

import java.io.File;
import java.io.FileInputStream;
Expand All @@ -70,6 +71,8 @@
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.notNullValue;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.lessThan;
import static org.junit.jupiter.params.provider.Arguments.arguments;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
Expand Down Expand Up @@ -324,6 +327,99 @@ void testWithDynamicGroups() throws IOException {
}
}

@Test
void testWithDynamicGroupsAndAggregateThreshold() throws IOException {
final BufferScenario bufferScenario = new InMemoryBufferScenario();
final CompressionScenario compressionScenario = new NoneCompressionScenario();
final NdjsonOutputScenario outputScenario = new NdjsonOutputScenario();
final SizeCombination sizeCombination = SizeCombination.MEDIUM_SMALLER;

BufferTypeOptions bufferTypeOptions = bufferScenario.getBufferType();
String testRun = "grouping-" + outputScenario + "-" + bufferTypeOptions + "-" + compressionScenario + "-" + sizeCombination.getBatchSize() + "-" + sizeCombination.getNumberOfBatches();

final String staticPrefix = "s3-sink-grouping-integration-test/aggregate-threshold-" + UUID.randomUUID() + "/";
final String pathPrefix = staticPrefix + "folder-${/sequence}/";
final List<String> dynamicKeys = new ArrayList<>();
dynamicKeys.add("/sequence");

final AggregateThresholdOptions aggregateThresholdOptions = mock(AggregateThresholdOptions.class);
when(aggregateThresholdOptions.getMaximumSize()).thenReturn(ByteCount.parse("100kb"));
when(s3SinkConfig.getAggregateThresholdOptions()).thenReturn(aggregateThresholdOptions);

when(objectKeyOptions.getPathPrefix()).thenReturn(pathPrefix);

when(expressionEvaluator.extractDynamicExpressionsFromFormatExpression(objectKeyOptions.getPathPrefix()))
.thenReturn(Collections.emptyList());
when(expressionEvaluator.extractDynamicKeysFromFormatExpression(objectKeyOptions.getPathPrefix()))
.thenReturn(dynamicKeys);
when(expressionEvaluator.extractDynamicKeysFromFormatExpression(objectKeyOptions.getNamePattern()))
.thenReturn(Collections.emptyList());
when(expressionEvaluator.extractDynamicExpressionsFromFormatExpression(objectKeyOptions.getNamePattern()))
.thenReturn(Collections.emptyList());

when(pluginFactory.loadPlugin(eq(OutputCodec.class), any())).thenReturn(outputScenario.getCodec());
when(s3SinkConfig.getBufferType()).thenReturn(bufferTypeOptions);
when(s3SinkConfig.getCompression()).thenReturn(compressionScenario.getCompressionOption());
int expectedTotalSize = sizeCombination.getTotalSize();

// This will not be hit since these are grouped dynamically
when(thresholdOptions.getEventCount()).thenReturn(expectedTotalSize);

final S3Sink objectUnderTest = createObjectUnderTest();

final int maxEventDataToSample = 2000;
final List<Map<String, Object>> sampleEventData = new ArrayList<>(maxEventDataToSample);
for (int batchNumber = 0; batchNumber < sizeCombination.getNumberOfBatches(); batchNumber++) {
final int currentBatchNumber = batchNumber;
final List<Record<Event>> events = IntStream.range(0, sizeCombination.getBatchSize())
.mapToObj(this::generateEventData)
.peek(data -> {
if (sampleEventData.size() < maxEventDataToSample)
sampleEventData.add(data);
})
.map(this::generateTestEvent)
.map(Record::new)
.collect(Collectors.toList());

LOG.debug("Writing dynamic batch {} with size {}.", currentBatchNumber, events.size());
objectUnderTest.doOutput(events);
}

for (int folderNumber = 0; folderNumber < 100; folderNumber++) {
LOG.info("Listing S3 path prefix: {}", staticPrefix + "folder-" + folderNumber + "/");

final ListObjectsV2Response listObjectsResponse = s3Client.listObjectsV2(ListObjectsV2Request.builder()
.bucket(bucketName)
.prefix(staticPrefix + "folder-" + folderNumber + "/")
.build());

assertThat(listObjectsResponse.contents(), notNullValue());
assertThat(listObjectsResponse.contents().size(), greaterThanOrEqualTo(1));

int objectIndex = 0;
for (final S3Object s3Object : listObjectsResponse.contents()) {
final File target = new File(s3FileLocation, "folder-" + folderNumber + "-" + objectIndex + ".original");

LOG.info("Downloading S3 object to local file {}.", target);

GetObjectRequest getObjectRequest = GetObjectRequest.builder()
.bucket(bucketName)
.key(s3Object.key())
.build();
s3Client.getObject(getObjectRequest, target.toPath());

File actualContentFile = decompressFileIfNecessary(outputScenario, compressionScenario, testRun, target);

// Since data is evenly distributed between all groups, and flush_capacity_ratio is 0
// all objects should be less than 10 kb (100 groups * 10 kb = aggregate threshold of 1 mb)
assertThat(actualContentFile.length(), lessThan(10_000L));
LOG.info("Validating output. object size is {}", actualContentFile.length());
outputScenario.validateDynamicPartition(-1, folderNumber, actualContentFile, compressionScenario);
objectIndex++;
}
}
}

private File decompressFileIfNecessary(OutputScenario outputScenario, CompressionScenario compressionScenario, String pathPrefix, File target) throws IOException {

if (outputScenario.isCompressionInternal() || !compressionScenario.requiresDecompression())
Expand Down
14 changes: 13 additions & 1 deletion ...lugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import org.opensearch.dataprepper.model.configuration.PluginModel;
import org.opensearch.dataprepper.plugins.sink.s3.accumulator.BufferTypeOptions;
import org.opensearch.dataprepper.plugins.sink.s3.compression.CompressionOption;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.AggregateThresholdOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.AwsAuthenticationOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.ObjectKeyOptions;
import org.opensearch.dataprepper.plugins.sink.s3.configuration.ThresholdOptions;
Expand All @@ -37,15 +38,21 @@ public class S3SinkConfig {
private String bucketName;

@JsonProperty("object_key")
@Valid
private ObjectKeyOptions objectKeyOptions = new ObjectKeyOptions();

@JsonProperty("compression")
private CompressionOption compression = CompressionOption.NONE;

@JsonProperty("threshold")
@NotNull
@Valid
private ThresholdOptions thresholdOptions;

@JsonProperty("aggregate_threshold")
@Valid
private AggregateThresholdOptions aggregateThresholdOptions;
Copy link
Member Author

Choose a reason for hiding this comment

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

Should this be made to private AggregateThresholdOptions aggregateThresholdOptions = new AggregateThresholdOptions();, or should this remain null by default?

Copy link
Member

@dlvenable dlvenable Apr 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

I like the idea of having a default. But, I think it would have to by dynamic. That is, we'd probably need to assign it like this:

max(DEFAULT_BYTE_CAPACITY, threshold.maximum_size)

As it is now, we may have a bug if you set this:

aggregate_threshold:
  flush_capacity_ratio: 0.25
threshold:
  maximum_size: 100mb

You would have an aggregate threshold of 50mb, but you are trying to get a 100mb threshold on each group.

Copy link
Member Author

Choose a reason for hiding this comment

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

This makes sense. It is a little strange to add defaults based on other config parameters though, as they aren't really connected in the code. Also, it may be better to not require the code that checks for aggregate thresholds to run every time by default, since in many cases (all existing cases) users only have a single group, making the aggregate threshold obsolete. For these reasons, I will just leave it with no default

Copy link
Member

Choose a reason for hiding this comment

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

Yea, let's not have any default - even the maximum_size. I think if the user is going to specify the aggregate_threshold at all, then he must also specify the maximum_size. This will require the user to think about it more and hopefully get it correct in relation to the other thresholds.

We could always change this later. But, if you remove the default 50mb for now, we won't have this possible confusion.


@JsonProperty("codec")
@NotNull
private PluginModel codec;
Expand All @@ -67,13 +74,18 @@ public AwsAuthenticationOptions getAwsAuthenticationOptions() {
}

/**
* Threshold configuration Options.
* Threshold configuration Options at the individual S3 group level
* @return threshold option object.
*/
public ThresholdOptions getThresholdOptions() {
return thresholdOptions;
}

/**
* Threshold configuration for the aggregation of all S3 groups
*/
public AggregateThresholdOptions getAggregateThresholdOptions() { return aggregateThresholdOptions; }

/**
* Read s3 bucket name configuration.
* @return bucket name.
Expand Down
65 changes: 53 additions & 12 deletions ...ugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@ public class S3SinkService {
public static final String OBJECTS_FAILED = "s3SinkObjectsFailed";
public static final String NUMBER_OF_RECORDS_FLUSHED_TO_S3_SUCCESS = "s3SinkObjectsEventsSucceeded";
public static final String NUMBER_OF_RECORDS_FLUSHED_TO_S3_FAILED = "s3SinkObjectsEventsFailed";

private static final String CURRENT_S3_GROUPS = "s3SinkNumberOfGroups";

static final String NUMBER_OF_GROUPS_FORCE_FLUSHED = "s3SinkObjectsForceFlushed";
static final String S3_OBJECTS_SIZE = "s3SinkObjectSizeBytes";
private final S3SinkConfig s3SinkConfig;
private final Lock reentrantLock;
Expand All @@ -56,6 +60,8 @@ public class S3SinkService {
private final Counter numberOfRecordsSuccessCounter;
private final Counter numberOfRecordsFailedCounter;
private final DistributionSummary s3ObjectSizeSummary;

private final Counter numberOfObjectsForceFlushed;
private final OutputCodecContext codecContext;
private final KeyGenerator keyGenerator;
private final Duration retrySleepTime;
Expand Down Expand Up @@ -91,6 +97,9 @@ public S3SinkService(final S3SinkConfig s3SinkConfig, final OutputCodec codec,
numberOfRecordsSuccessCounter = pluginMetrics.counter(NUMBER_OF_RECORDS_FLUSHED_TO_S3_SUCCESS);
numberOfRecordsFailedCounter = pluginMetrics.counter(NUMBER_OF_RECORDS_FLUSHED_TO_S3_FAILED);
s3ObjectSizeSummary = pluginMetrics.summary(S3_OBJECTS_SIZE);
numberOfObjectsForceFlushed = pluginMetrics.counter(NUMBER_OF_GROUPS_FORCE_FLUSHED);
pluginMetrics.gauge(CURRENT_S3_GROUPS, s3GroupManager, S3GroupManager::getNumberOfGroups);


this.s3GroupManager = s3GroupManager;
}
Expand All @@ -110,10 +119,10 @@ void output(Collection<Record<Event>> records) {
try {
for (Record<Event> record : records) {
final Event event = record.getData();
final S3Group s3Group = s3GroupManager.getOrCreateGroupForEvent(event);
final Buffer currentBuffer = s3Group.getBuffer();

try {
final S3Group s3Group = s3GroupManager.getOrCreateGroupForEvent(event);
final Buffer currentBuffer = s3Group.getBuffer();

if (currentBuffer.getEventCount() == 0) {
codec.start(currentBuffer.getOutputStream(), event, codecContext);
}
Expand All @@ -122,28 +131,32 @@ void output(Collection<Record<Event>> records) {
int count = currentBuffer.getEventCount() + 1;
currentBuffer.setEventCount(count);
s3Group.addEventHandle(event.getEventHandle());

final boolean flushed = flushToS3IfNeeded(s3Group, false);

if (flushed) {
s3GroupManager.removeGroup(s3Group);
}
} catch (Exception ex) {
if(sampleException == null) {
sampleException = ex;
}

failedEvents.add(event);
}

final boolean flushed = flushToS3IfNeeded(s3Group);

if (flushed) {
s3GroupManager.removeGroup(s3Group);
}
}

for (final S3Group s3Group : s3GroupManager.getS3GroupEntries()) {
final boolean flushed = flushToS3IfNeeded(s3Group);
final boolean flushed = flushToS3IfNeeded(s3Group, false);

if (flushed) {
s3GroupManager.removeGroup(s3Group);
}
}

if (s3SinkConfig.getAggregateThresholdOptions() != null) {
checkAggregateThresholdsAndFlushIfNeeded();
}
} finally {
reentrantLock.unlock();
}
Expand All @@ -160,10 +173,10 @@ void output(Collection<Record<Event>> records) {
/**
* @return whether the flush was attempted
*/
private boolean flushToS3IfNeeded(final S3Group s3Group) {
private boolean flushToS3IfNeeded(final S3Group s3Group, final boolean forceFlush) {
LOG.trace("Flush to S3 check: currentBuffer.size={}, currentBuffer.events={}, currentBuffer.duration={}",
s3Group.getBuffer().getSize(), s3Group.getBuffer().getEventCount(), s3Group.getBuffer().getDuration());
if (ThresholdCheck.checkThresholdExceed(s3Group.getBuffer(), maxEvents, maxBytes, maxCollectionDuration)) {
if (forceFlush || ThresholdCheck.checkThresholdExceed(s3Group.getBuffer(), maxEvents, maxBytes, maxCollectionDuration)) {
try {
codec.complete(s3Group.getBuffer().getOutputStream());
String s3Key = s3Group.getBuffer().getKey();
Expand Down Expand Up @@ -224,4 +237,32 @@ protected boolean retryFlushToS3(final Buffer currentBuffer, final String s3Key)
} while (!isUploadedToS3);
return isUploadedToS3;
}

private void checkAggregateThresholdsAndFlushIfNeeded() {
long currentTotalGroupSize = s3GroupManager.recalculateAndGetGroupSize();
LOG.debug("Total groups size is {} bytes", currentTotalGroupSize);

final long aggregateThresholdBytes = s3SinkConfig.getAggregateThresholdOptions().getMaximumSize().getBytes();
final double aggregateThresholdFlushRatio = s3SinkConfig.getAggregateThresholdOptions().getFlushCapacityRatio();

if (currentTotalGroupSize >= aggregateThresholdBytes) {
LOG.info("aggregate_threshold reached, the largest groups will be flushed until {} percent of the maximum size {} is remaining", aggregateThresholdFlushRatio * 100, aggregateThresholdBytes);

for (final S3Group s3Group : s3GroupManager.getS3GroupsSortedBySize()) {
LOG.info("Forcing a flush of object with key {} due to aggregate_threshold of {} bytes being reached", s3Group.getBuffer().getKey(), aggregateThresholdBytes);

boolean flushed = flushToS3IfNeeded(s3Group, true);
numberOfObjectsForceFlushed.increment();

if (flushed) {
currentTotalGroupSize -= s3Group.getBuffer().getSize();
s3GroupManager.removeGroup(s3Group);
}

if (currentTotalGroupSize <= aggregateThresholdBytes * aggregateThresholdFlushRatio) {
break;
}
}
}
}
}
42 changes: 42 additions & 0 deletions ...a/org/opensearch/dataprepper/plugins/sink/s3/configuration/AggregateThresholdOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.dataprepper.plugins.sink.s3.configuration;

import com.fasterxml.jackson.annotation.JsonProperty;
import jakarta.validation.constraints.Max;
import jakarta.validation.constraints.Min;
import jakarta.validation.constraints.NotNull;
import org.opensearch.dataprepper.model.types.ByteCount;


/**
* An implementation class of s3 index configuration Options
*/
public class AggregateThresholdOptions {

@JsonProperty("maximum_size")
@NotNull
private ByteCount maximumSize;

/**
* Controls how aggressive the force flush is when the maximum_size is reached.
* Groups will be flushed until the total size is less than maximum_size * flush_capacity_ratio
*/
@JsonProperty("flush_capacity_ratio")
@Min(value = 0, message = "flush_capacity_ratio must be between 0.0 and 1.0")
@Max(value = 1, message = "flush_capacity_ratio must be between 0.0 and 1.0")
private double flushCapacityRatio = 0.5;

/**
* Read byte capacity configuration.
* @return maximum byte count.
*/
public ByteCount getMaximumSize() {
return maximumSize;
}

public double getFlushCapacityRatio() { return flushCapacityRatio; }
}
10 changes: 6 additions & 4 deletions .../main/java/org/opensearch/dataprepper/plugins/sink/s3/configuration/ThresholdOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,14 @@
package org.opensearch.dataprepper.plugins.sink.s3.configuration;

import java.time.Duration;

import jakarta.validation.constraints.Max;
import jakarta.validation.constraints.Min;
import jakarta.validation.constraints.NotNull;
import org.hibernate.validator.constraints.time.DurationMax;
import org.hibernate.validator.constraints.time.DurationMin;
import org.opensearch.dataprepper.model.types.ByteCount;
import com.fasterxml.jackson.annotation.JsonProperty;
import jakarta.validation.constraints.NotNull;
import jakarta.validation.constraints.Size;

/**
* An implementation class of s3 index configuration Options
Expand All @@ -21,8 +23,8 @@ public class ThresholdOptions {
private static final String DEFAULT_BYTE_CAPACITY = "50mb";

@JsonProperty("event_count")
@Size(min = 0, max = 10000000, message = "event_count size should be between 0 and 10000000")
@NotNull
@Min(value = 0, message = "event_count size should be between 0 and 10000000")
@Max(value = 10000000, message = "event_count size should be between 0 and 10000000")
private int eventCount;

@JsonProperty("maximum_size")
Expand Down
Loading
Loading