Skip to content

Commit aebb597

Browse files
graytaylor0graytaylor0
authored
Fix issue where event handles are not released for cloudwatch logs sink, add more tests (#6074)
Signed-off-by: Taylor Gray <tylgry@amazon.com>
1 parent 9b81626 commit aebb597

2 files changed

Lines changed: 193 additions & 4 deletions

File tree

‎data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsDispatcher.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ public void upload() {
164164
dlqObjects = getDlqObjectsFromResponse(putLogEventsResponse);
165165
}
166166
cloudWatchLogsMetrics.increaseLogEventSuccessCounter(totalEventCount - dlqObjects.size());
167+
releaseEventHandles(putLogEventsResponse);
167168
}
168169
CloudWatchLogsSinkUtils.handleDlqObjects(dlqObjects, dlqPushHandler);
169170
}
@@ -200,5 +201,23 @@ List<DlqObject> getDlqObjectsFromResponse(PutLogEventsResponse putLogEventsRespo
200201
return dlqObjects;
201202
}
202203

204+
private void releaseEventHandles(final PutLogEventsResponse putLogEventsResponse) {
205+
if (putLogEventsResponse == null || putLogEventsResponse.rejectedLogEventsInfo() == null) {
206+
eventHandles.forEach(eventHandle -> eventHandle.release(true));
207+
return;
208+
}
209+
210+
final Integer tooOldEndIndex = putLogEventsResponse.rejectedLogEventsInfo().tooOldLogEventEndIndex();
211+
final Integer tooNewStartIndex = putLogEventsResponse.rejectedLogEventsInfo().tooNewLogEventStartIndex();
212+
213+
for (int i = 0; i < eventHandles.size(); i++) {
214+
boolean isRejected = (tooOldEndIndex != null && i < tooOldEndIndex) ||
215+
(tooNewStartIndex != null && i >= tooNewStartIndex);
216+
217+
if (!isRejected) {
218+
eventHandles.get(i).release(true);
219+
}
220+
}
221+
}
203222
}
204223
}

‎data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsDispatcherTest.java‎

Lines changed: 174 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,32 @@
77

88
import org.junit.jupiter.api.BeforeEach;
99
import org.junit.jupiter.api.Test;
10+
import org.mockito.ArgumentCaptor;
1011
import org.opensearch.dataprepper.model.event.EventHandle;
1112
import org.opensearch.dataprepper.plugins.sink.cloudwatch_logs.config.ThresholdConfig;
1213
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
1314
import software.amazon.awssdk.services.cloudwatchlogs.model.InputLogEvent;
15+
import software.amazon.awssdk.services.cloudwatchlogs.model.PutLogEventsRequest;
16+
import software.amazon.awssdk.services.cloudwatchlogs.model.PutLogEventsResponse;
17+
import software.amazon.awssdk.services.cloudwatchlogs.model.RejectedLogEventsInfo;
18+
import software.amazon.awssdk.core.exception.SdkClientException;
19+
import software.amazon.awssdk.services.cloudwatchlogs.model.CloudWatchLogsException;
1420

1521
import java.util.ArrayList;
1622
import java.util.Collection;
1723
import java.util.List;
1824
import java.util.concurrent.Executor;
1925

2026
import static org.mockito.ArgumentMatchers.any;
21-
import static org.mockito.Mockito.atMostOnce;
2227
import static org.mockito.Mockito.mock;
2328
import static org.mockito.Mockito.verify;
29+
import static org.mockito.Mockito.when;
30+
import static org.mockito.Mockito.times;
31+
import static org.mockito.Mockito.never;
2432

2533
class CloudWatchLogsDispatcherTest {
2634
private CloudWatchLogsDispatcher cloudWatchLogsDispatcher;
27-
private CloudWatchLogsClient mockCloudWatchLogsClient;
35+
private CloudWatchLogsClient mockCloudWatchLogsClient;
2836
private CloudWatchLogsMetrics mockCloudWatchLogsMetrics;
2937
private Executor mockExecutor;
3038
private static final int RETRY_COUNT = 5;
@@ -71,12 +79,174 @@ CloudWatchLogsDispatcher getCloudWatchLogsDispatcher() {
7179
.build();
7280
}
7381

82+
private void executeDispatcherRunnable() {
83+
ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
84+
verify(mockExecutor).execute(runnableCaptor.capture());
85+
runnableCaptor.getValue().run();
86+
}
87+
7488
@Test
7589
void GIVEN_valid_input_log_events_SHOULD_call_executor() {
7690
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
91+
92+
final List<EventHandle> eventHandles = getSampleEventHandles();
93+
final PutLogEventsResponse response = mock(PutLogEventsResponse.class);
94+
when(response.rejectedLogEventsInfo()).thenReturn(null);
95+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class))).thenReturn(response);
96+
97+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
98+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
99+
100+
executeDispatcherRunnable();
101+
102+
eventHandles.forEach(eventHandle -> verify(eventHandle).release(true));
103+
}
104+
105+
@Test
106+
void GIVEN_too_old_events_SHOULD_not_release_old_events() {
107+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
108+
109+
final List<EventHandle> eventHandles = getSampleEventHandles();
110+
final PutLogEventsResponse response = mock(PutLogEventsResponse.class);
111+
final RejectedLogEventsInfo rejectedInfo = mock(RejectedLogEventsInfo.class);
112+
113+
// Set first 2 events as too old
114+
when(rejectedInfo.tooOldLogEventEndIndex()).thenReturn(2);
115+
when(rejectedInfo.tooNewLogEventStartIndex()).thenReturn(null);
116+
when(response.rejectedLogEventsInfo()).thenReturn(rejectedInfo);
117+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class))).thenReturn(response);
118+
119+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
120+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
121+
122+
executeDispatcherRunnable();
123+
124+
// First two events should not be released
125+
verify(eventHandles.get(0), never()).release(true);
126+
verify(eventHandles.get(1), never()).release(true);
127+
128+
// Remaining events should be released
129+
for (int i = 2; i < eventHandles.size(); i++) {
130+
verify(eventHandles.get(i)).release(true);
131+
}
132+
}
133+
134+
@Test
135+
void GIVEN_too_new_events_SHOULD_not_release_new_events() {
136+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
137+
138+
final List<EventHandle> eventHandles = getSampleEventHandles();
139+
final PutLogEventsResponse response = mock(PutLogEventsResponse.class);
140+
final RejectedLogEventsInfo rejectedInfo = mock(RejectedLogEventsInfo.class);
141+
142+
// Set last 3 events as too new
143+
when(rejectedInfo.tooOldLogEventEndIndex()).thenReturn(null);
144+
when(rejectedInfo.tooNewLogEventStartIndex()).thenReturn(eventHandles.size() - 3);
145+
when(response.rejectedLogEventsInfo()).thenReturn(rejectedInfo);
146+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class))).thenReturn(response);
147+
77148
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
78-
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, getSampleEventHandles());
149+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
150+
151+
executeDispatcherRunnable();
152+
153+
// Events before tooNewStartIndex should be released
154+
for (int i = 0; i < eventHandles.size() - 3; i++) {
155+
verify(eventHandles.get(i)).release(true);
156+
}
157+
158+
// Last three events should not be released
159+
for (int i = eventHandles.size() - 3; i < eventHandles.size(); i++) {
160+
verify(eventHandles.get(i), never()).release(true);
161+
}
162+
}
163+
164+
@Test
165+
void GIVEN_both_old_and_new_rejected_events_SHOULD_only_release_valid_events() {
166+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
167+
168+
final List<EventHandle> eventHandles = getSampleEventHandles();
169+
final PutLogEventsResponse response = mock(PutLogEventsResponse.class);
170+
final RejectedLogEventsInfo rejectedInfo = mock(RejectedLogEventsInfo.class);
171+
172+
// Set first 2 events as too old and last 2 as too new
173+
when(rejectedInfo.tooOldLogEventEndIndex()).thenReturn(2);
174+
when(rejectedInfo.tooNewLogEventStartIndex()).thenReturn(eventHandles.size() - 2);
175+
when(response.rejectedLogEventsInfo()).thenReturn(rejectedInfo);
176+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class))).thenReturn(response);
177+
178+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
179+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
180+
181+
executeDispatcherRunnable();
182+
183+
// First two events should not be released (too old)
184+
verify(eventHandles.get(0), never()).release(true);
185+
verify(eventHandles.get(1), never()).release(true);
186+
187+
// Middle events should be released
188+
for (int i = 2; i < eventHandles.size() - 2; i++) {
189+
verify(eventHandles.get(i)).release(true);
190+
}
191+
192+
// Last two events should not be released (too new)
193+
verify(eventHandles.get(eventHandles.size() - 2), never()).release(true);
194+
verify(eventHandles.get(eventHandles.size() - 1), never()).release(true);
195+
}
196+
197+
@Test
198+
void GIVEN_client_exception_SHOULD_retry() {
199+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
200+
201+
final List<EventHandle> eventHandles = getSampleEventHandles();
202+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class)))
203+
.thenThrow(SdkClientException.create("Test exception"))
204+
.thenReturn(mock(PutLogEventsResponse.class));
205+
206+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
207+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
208+
209+
executeDispatcherRunnable();
210+
211+
verify(mockCloudWatchLogsMetrics, times(1)).increaseRequestFailCounter(1);
212+
verify(mockCloudWatchLogsMetrics, times(1)).increaseRequestSuccessCounter(1);
213+
}
214+
215+
@Test
216+
void GIVEN_cloudwatch_exception_SHOULD_retry() {
217+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
218+
219+
final List<EventHandle> eventHandles = getSampleEventHandles();
220+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class)))
221+
.thenThrow(CloudWatchLogsException.class)
222+
.thenReturn(mock(PutLogEventsResponse.class));
223+
224+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
225+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
226+
227+
executeDispatcherRunnable();
228+
229+
verify(mockCloudWatchLogsMetrics, times(1)).increaseRequestFailCounter(1);
230+
verify(mockCloudWatchLogsMetrics, times(1)).increaseRequestSuccessCounter(1);
231+
}
232+
233+
@Test
234+
void GIVEN_max_retries_exceeded_SHOULD_not_release_events() {
235+
cloudWatchLogsDispatcher = getCloudWatchLogsDispatcher();
236+
237+
final List<EventHandle> eventHandles = getSampleEventHandles();
238+
when(mockCloudWatchLogsClient.putLogEvents(any(PutLogEventsRequest.class)))
239+
.thenThrow(CloudWatchLogsException.class);
240+
241+
List<InputLogEvent> inputLogEventList = cloudWatchLogsDispatcher.prepareInputLogEvents(getSampleBufferedData());
242+
cloudWatchLogsDispatcher.dispatchLogs(inputLogEventList, eventHandles);
243+
244+
executeDispatcherRunnable();
245+
246+
verify(mockCloudWatchLogsMetrics, times(RETRY_COUNT)).increaseRequestFailCounter(1);
247+
verify(mockCloudWatchLogsMetrics, never()).increaseRequestSuccessCounter(1);
79248

80-
verify(mockExecutor, atMostOnce()).execute(any(CloudWatchLogsDispatcher.Uploader.class));
249+
// No events should be released after max retries
250+
eventHandles.forEach(eventHandle -> verify(eventHandle, never()).release(true));
81251
}
82252
}

0 commit comments

Comments
 (0)