|
9 | 9 |
|
10 | 10 | package org.elasticsearch.ingest; |
11 | 11 |
|
| 12 | +import org.elasticsearch.cluster.metadata.DataStream; |
12 | 13 | import org.elasticsearch.common.bytes.BytesArray; |
13 | 14 | import org.elasticsearch.common.xcontent.XContentHelper; |
14 | 15 | import org.elasticsearch.test.ESTestCase; |
15 | 16 | import org.elasticsearch.xcontent.XContentType; |
16 | 17 | import org.hamcrest.Matchers; |
| 18 | +import org.junit.Assume; |
17 | 19 | import org.junit.Before; |
| 20 | +import org.mockito.ArgumentCaptor; |
18 | 21 |
|
19 | 22 | import java.time.Instant; |
20 | 23 | import java.time.ZoneId; |
|
25 | 28 | import java.util.List; |
26 | 29 | import java.util.Map; |
27 | 30 | import java.util.Set; |
| 31 | +import java.util.function.BiConsumer; |
28 | 32 | import java.util.stream.DoubleStream; |
29 | 33 |
|
30 | 34 | import static org.elasticsearch.ingest.IngestDocumentMatcher.assertIngestDocument; |
31 | 35 | import static org.hamcrest.Matchers.both; |
32 | 36 | import static org.hamcrest.Matchers.containsInAnyOrder; |
33 | 37 | import static org.hamcrest.Matchers.containsString; |
| 38 | +import static org.hamcrest.Matchers.empty; |
34 | 39 | import static org.hamcrest.Matchers.equalTo; |
35 | 40 | import static org.hamcrest.Matchers.greaterThanOrEqualTo; |
36 | 41 | import static org.hamcrest.Matchers.instanceOf; |
|
40 | 45 | import static org.hamcrest.Matchers.notNullValue; |
41 | 46 | import static org.hamcrest.Matchers.nullValue; |
42 | 47 | import static org.hamcrest.Matchers.sameInstance; |
| 48 | +import static org.mockito.ArgumentMatchers.eq; |
| 49 | +import static org.mockito.Mockito.mock; |
| 50 | +import static org.mockito.Mockito.verify; |
| 51 | +import static org.mockito.Mockito.when; |
43 | 52 |
|
44 | 53 | public class IngestDocumentTests extends ESTestCase { |
45 | 54 |
|
@@ -1245,4 +1254,95 @@ public void testSourceHashMapIsNotCopied() { |
1245 | 1254 | assertThat(document2.getCtxMap().getMetadata(), not(sameInstance(document1.getCtxMap().getMetadata()))); |
1246 | 1255 | } |
1247 | 1256 | } |
| 1257 | + |
| 1258 | + /** |
| 1259 | + * When executing nested pipelines on an ingest document, the document should keep track of each pipeline's access pattern for the |
| 1260 | + * lifetime of each pipeline execution. When a pipeline execution concludes, it should clear access pattern from the document and |
| 1261 | + * restore the previous pipeline's access pattern. |
| 1262 | + */ |
| 1263 | + public void testNestedAccessPatternPropagation() { |
| 1264 | + Assume.assumeTrue(DataStream.LOGS_STREAM_FEATURE_FLAG); |
| 1265 | + |
| 1266 | + Map<String, Object> source = new HashMap<>(Map.of("foo", 1)); |
| 1267 | + IngestDocument document = new IngestDocument("index", "id", 1, null, null, source); |
| 1268 | + |
| 1269 | + // 1-3 nested calls |
| 1270 | + doTestNestedAccessPatternPropagation(0, randomIntBetween(1, 5), document); |
| 1271 | + |
| 1272 | + // At the end of the test, there should be neither pipeline ids nor access patterns left in the stack. |
| 1273 | + assertThat(document.getPipelineStack(), is(empty())); |
| 1274 | + assertThat(document.getCurrentAccessPattern(), is(nullValue())); |
| 1275 | + } |
| 1276 | + |
| 1277 | + /** |
| 1278 | + * Recursively execute some number of pipelines at various call depths to simulate a robust chain of pipelines being called on a |
| 1279 | + * document. |
| 1280 | + * @param level The current call depth. This is how many pipelines deep into the nesting we are. |
| 1281 | + * @param maxCallDepth How much further in the call depth we should go in the test. If this is greater than the current level, we will |
| 1282 | + * recurse in at least one of the pipelines executed at this level. If the current level is equal to the max call |
| 1283 | + * depth we will run some pipelines but recurse no further before returning. |
| 1284 | + * @param document The document to repeatedly use and verify against. |
| 1285 | + */ |
| 1286 | + void doTestNestedAccessPatternPropagation(int level, int maxCallDepth, IngestDocument document) { |
| 1287 | + // 1-5 pipelines to be run at any given level |
| 1288 | + logger.debug("LEVEL {}/{}: BEGIN", level, maxCallDepth); |
| 1289 | + int pipelinesAtThisLevel = randomIntBetween(1, 7); |
| 1290 | + logger.debug("Run pipelines: {}", pipelinesAtThisLevel); |
| 1291 | + |
| 1292 | + boolean recursed = false; |
| 1293 | + if (level >= maxCallDepth) { |
| 1294 | + // If we're at max call depth, do no recursions |
| 1295 | + recursed = true; |
| 1296 | + logger.debug("No more recursions"); |
| 1297 | + } |
| 1298 | + |
| 1299 | + for (int pipelineIdx = 0; pipelineIdx < pipelinesAtThisLevel; pipelineIdx++) { |
| 1300 | + String expectedPipelineId = randomAlphaOfLength(20); |
| 1301 | + IngestPipelineFieldAccessPattern expectedAccessPattern = randomFrom(IngestPipelineFieldAccessPattern.values()); |
| 1302 | + |
| 1303 | + // We mock the pipeline because it's easier to verify calls and doesn't |
| 1304 | + // need us to force a stall in the execution logic to half apply it. |
| 1305 | + Pipeline mockPipeline = mock(Pipeline.class); |
| 1306 | + when(mockPipeline.getId()).thenReturn(expectedPipelineId); |
| 1307 | + when(mockPipeline.getProcessors()).thenReturn(List.of(new TestProcessor((doc) -> {}))); |
| 1308 | + when(mockPipeline.getFieldAccessPattern()).thenReturn(expectedAccessPattern); |
| 1309 | + @SuppressWarnings("unchecked") |
| 1310 | + BiConsumer<IngestDocument, Exception> mockHandler = mock(BiConsumer.class); |
| 1311 | + |
| 1312 | + // Execute pipeline |
| 1313 | + logger.debug("LEVEL {}/{}: Executing {}/{}", level, maxCallDepth, pipelineIdx, pipelinesAtThisLevel); |
| 1314 | + document.executePipeline(mockPipeline, mockHandler); |
| 1315 | + |
| 1316 | + // Verify pipeline was called, capture completion handler |
| 1317 | + ArgumentCaptor<BiConsumer<IngestDocument, Exception>> argumentCaptor = ArgumentCaptor.captor(); |
| 1318 | + verify(mockPipeline).execute(eq(document), argumentCaptor.capture()); |
| 1319 | + |
| 1320 | + // Assert expected state |
| 1321 | + assertThat(document.getPipelineStack().getFirst(), is(expectedPipelineId)); |
| 1322 | + assertThat(document.getCurrentAccessPattern(), is(expectedAccessPattern)); |
| 1323 | + |
| 1324 | + // Randomly recurse: We recurse only one time per level to avoid hogging test time, but we randomize which |
| 1325 | + // pipeline to recurse on, eventually requiring a recursion on the last pipeline run if one hasn't happened yet. |
| 1326 | + if (recursed == false && (randomBoolean() || pipelineIdx == pipelinesAtThisLevel - 1)) { |
| 1327 | + logger.debug("Recursed on pipeline {}", pipelineIdx); |
| 1328 | + doTestNestedAccessPatternPropagation(level + 1, maxCallDepth, document); |
| 1329 | + recursed = true; |
| 1330 | + } |
| 1331 | + |
| 1332 | + // Pull up the captured completion handler to conclude the pipeline run |
| 1333 | + argumentCaptor.getValue().accept(document, null); |
| 1334 | + |
| 1335 | + // Assert expected state |
| 1336 | + assertThat(document.getPipelineStack().size(), is(equalTo(level))); |
| 1337 | + if (level == 0) { |
| 1338 | + // Top level means access pattern should be empty |
| 1339 | + assertThat(document.getCurrentAccessPattern(), is(nullValue())); |
| 1340 | + } else { |
| 1341 | + // If we're nested below the top level we should still have an access |
| 1342 | + // pattern on the document for the pipeline above us |
| 1343 | + assertThat(document.getCurrentAccessPattern(), is(not(nullValue()))); |
| 1344 | + } |
| 1345 | + } |
| 1346 | + logger.debug("LEVEL {}/{}: COMPLETE", level, maxCallDepth); |
| 1347 | + } |
1248 | 1348 | } |
0 commit comments