-
Notifications
You must be signed in to change notification settings - Fork 207
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Use buffer accumulator in opensearch source to backoff and retry
Signed-off-by: Taylor Gray <[email protected]>
- Loading branch information
1 parent
2f979d5
commit 168227d
Showing
9 changed files
with
536 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
130 changes: 130 additions & 0 deletions
130
...n/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/BufferAccumulator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,130 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.dataprepper.plugins.source.opensearch.worker; | ||
|
|
||
| import org.opensearch.dataprepper.model.buffer.Buffer; | ||
| import org.opensearch.dataprepper.model.record.Record; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| import java.time.Duration; | ||
| import java.util.ArrayList; | ||
| import java.util.Collection; | ||
| import java.util.Objects; | ||
| import java.util.concurrent.ExecutionException; | ||
| import java.util.concurrent.Executors; | ||
| import java.util.concurrent.ScheduledExecutorService; | ||
| import java.util.concurrent.ScheduledFuture; | ||
| import java.util.concurrent.TimeUnit; | ||
| import java.util.concurrent.TimeoutException; | ||
|
|
||
| /** | ||
| * Accumulates {@link Record} objects before placing them into a Data Prepper | ||
| * {@link Buffer}. This class is not thread-safe and should only be used by one | ||
| * thread at a time. | ||
| * | ||
| * @param <T> Type of record to accumulate | ||
| */ | ||
| public class BufferAccumulator<T extends Record<?>> { | ||
| private static final Logger LOG = LoggerFactory.getLogger(BufferAccumulator.class); | ||
|
|
||
| private static final int MAX_FLUSH_RETRIES_ON_IO_EXCEPTION = Integer.MAX_VALUE; | ||
| private static final Duration INITIAL_FLUSH_RETRY_DELAY_ON_IO_EXCEPTION = Duration.ofSeconds(5); | ||
|
|
||
| private final Buffer<T> buffer; | ||
| private final int numberOfRecordsToAccumulate; | ||
| private final int bufferTimeoutMillis; | ||
| private int totalWritten = 0; | ||
|
|
||
| private final Collection<T> recordsAccumulated; | ||
|
|
||
| private BufferAccumulator(final Buffer<T> buffer, final int numberOfRecordsToAccumulate, final Duration bufferTimeout) { | ||
| this.buffer = Objects.requireNonNull(buffer, "buffer must be non-null."); | ||
| this.numberOfRecordsToAccumulate = numberOfRecordsToAccumulate; | ||
| Objects.requireNonNull(bufferTimeout, "bufferTimeout must be non-null."); | ||
| this.bufferTimeoutMillis = (int) bufferTimeout.toMillis(); | ||
|
|
||
| if(numberOfRecordsToAccumulate < 1) | ||
| throw new IllegalArgumentException("numberOfRecordsToAccumulate must be greater than zero."); | ||
|
|
||
| recordsAccumulated = new ArrayList<>(numberOfRecordsToAccumulate); | ||
| } | ||
|
|
||
| public static <T extends Record<?>> BufferAccumulator<T> create(final Buffer<T> buffer, final int recordsToAccumulate, final Duration bufferTimeout) { | ||
| return new BufferAccumulator<T>(buffer, recordsToAccumulate, bufferTimeout); | ||
| } | ||
|
|
||
| void add(final T record) throws Exception { | ||
| recordsAccumulated.add(record); | ||
| if (recordsAccumulated.size() >= numberOfRecordsToAccumulate) { | ||
| flush(); | ||
| } | ||
| } | ||
|
|
||
| void flush() throws Exception { | ||
| try { | ||
| flushAccumulatedToBuffer(); | ||
| } catch (final TimeoutException timeoutException) { | ||
| flushWithBackoff(); | ||
| } | ||
| } | ||
|
|
||
| private boolean flushWithBackoff() throws Exception{ | ||
| final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); | ||
| long nextDelay = INITIAL_FLUSH_RETRY_DELAY_ON_IO_EXCEPTION.toMillis(); | ||
| boolean flushedSuccessfully; | ||
|
|
||
| for (int retryCount = 0; retryCount < MAX_FLUSH_RETRIES_ON_IO_EXCEPTION; retryCount++) { | ||
| final ScheduledFuture<Boolean> flushBufferFuture = scheduledExecutorService.schedule(() -> { | ||
| try { | ||
| flushAccumulatedToBuffer(); | ||
| return true; | ||
| } catch (final TimeoutException e) { | ||
| return false; | ||
| } | ||
| }, nextDelay, TimeUnit.MILLISECONDS); | ||
|
|
||
| try { | ||
| flushedSuccessfully = flushBufferFuture.get(); | ||
| if (flushedSuccessfully) { | ||
| LOG.info("Successfully flushed the buffer accumulator on retry attempt {}", retryCount + 1); | ||
| scheduledExecutorService.shutdownNow(); | ||
| return true; | ||
| } | ||
| } catch (final ExecutionException e) { | ||
| LOG.warn("Retrying of flushing the buffer accumulator hit an exception: {}", e.getMessage()); | ||
| scheduledExecutorService.shutdownNow(); | ||
| throw e; | ||
| } catch (final InterruptedException e) { | ||
| LOG.warn("Retrying of flushing the buffer accumulator was interrupted: {}", e.getMessage()); | ||
| scheduledExecutorService.shutdownNow(); | ||
| throw e; | ||
| } | ||
| } | ||
|
|
||
| LOG.warn("Flushing the bufferAccumulator failed after {} attempts", MAX_FLUSH_RETRIES_ON_IO_EXCEPTION); | ||
| scheduledExecutorService.shutdownNow(); | ||
| return false; | ||
| } | ||
|
|
||
| private void flushAccumulatedToBuffer() throws Exception { | ||
| final int currentRecordCountAccumulated = recordsAccumulated.size(); | ||
| if (currentRecordCountAccumulated > 0) { | ||
| buffer.writeAll(recordsAccumulated, bufferTimeoutMillis); | ||
| recordsAccumulated.clear(); | ||
| totalWritten += currentRecordCountAccumulated; | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Gets the total number of records written to the buffer. | ||
| * | ||
| * @return the total number of records written | ||
| */ | ||
| public int getTotalWritten() { | ||
| return totalWritten; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.