|  | 
| 9 | 9 | package org.opensearch.plugin.ingestion.fs; | 
| 10 | 10 | 
 | 
| 11 | 11 | import org.opensearch.action.admin.indices.delete.DeleteIndexRequest; | 
|  | 12 | +import org.opensearch.action.admin.indices.stats.IndexStats; | 
|  | 13 | +import org.opensearch.action.admin.indices.stats.ShardStats; | 
| 12 | 14 | import org.opensearch.action.admin.indices.streamingingestion.pause.PauseIngestionResponse; | 
| 13 | 15 | import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionRequest; | 
| 14 | 16 | import org.opensearch.action.admin.indices.streamingingestion.resume.ResumeIngestionResponse; | 
|  | 
| 17 | 19 | import org.opensearch.cluster.metadata.IndexMetadata; | 
| 18 | 20 | import org.opensearch.common.settings.Settings; | 
| 19 | 21 | import org.opensearch.index.query.RangeQueryBuilder; | 
|  | 22 | +import org.opensearch.indices.pollingingest.PollingIngestStats; | 
| 20 | 23 | import org.opensearch.plugins.Plugin; | 
| 21 | 24 | import org.opensearch.test.OpenSearchSingleNodeTestCase; | 
| 22 | 25 | import org.opensearch.transport.client.Requests; | 
|  | 
| 31 | 34 | import java.util.Arrays; | 
| 32 | 35 | import java.util.Collection; | 
| 33 | 36 | import java.util.Collections; | 
|  | 37 | +import java.util.concurrent.Callable; | 
|  | 38 | +import java.util.concurrent.TimeUnit; | 
| 34 | 39 | 
 | 
| 35 | 40 | public class FileBasedIngestionSingleNodeTests extends OpenSearchSingleNodeTestCase { | 
| 36 | 41 |     private Path ingestionDir; | 
| @@ -237,4 +242,165 @@ public void testFileIngestionFromProvidedPointer() throws Exception { | 
| 237 | 242 |         // cleanup the test index | 
| 238 | 243 |         client().admin().indices().delete(new DeleteIndexRequest(index)).actionGet(); | 
| 239 | 244 |     } | 
|  | 245 | + | 
|  | 246 | +    public void testPointerBasedLag() throws Exception { | 
|  | 247 | +        String mappings = """ | 
|  | 248 | +            { | 
|  | 249 | +              "properties": { | 
|  | 250 | +                "name": { "type": "text" }, | 
|  | 251 | +                "age": { "type": "integer" } | 
|  | 252 | +              } | 
|  | 253 | +            } | 
|  | 254 | +            """; | 
|  | 255 | + | 
|  | 256 | +        // Create index with empty file (no messages) | 
|  | 257 | +        Path streamDir = ingestionDir.resolve(stream); | 
|  | 258 | +        Path shardFile = streamDir.resolve("0.ndjson"); | 
|  | 259 | +        Files.write(shardFile, new byte[0]); // Empty file | 
|  | 260 | + | 
|  | 261 | +        createIndexWithMappingSource( | 
|  | 262 | +            index, | 
|  | 263 | +            Settings.builder() | 
|  | 264 | +                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) | 
|  | 265 | +                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) | 
|  | 266 | +                .put("ingestion_source.type", "FILE") | 
|  | 267 | +                .put("ingestion_source.pointer.init.reset", "earliest") | 
|  | 268 | +                .put("ingestion_source.param.stream", stream) | 
|  | 269 | +                .put("ingestion_source.param.base_directory", ingestionDir.toString()) | 
|  | 270 | +                .put("index.replication.type", "SEGMENT") | 
|  | 271 | +                .build(), | 
|  | 272 | +            mappings | 
|  | 273 | +        ); | 
|  | 274 | +        ensureGreen(index); | 
|  | 275 | + | 
|  | 276 | +        // Lag should be 0 since there are no messages | 
|  | 277 | +        waitForState(() -> { | 
|  | 278 | +            PollingIngestStats stats = getPollingIngestStats(index); | 
|  | 279 | +            return stats != null && stats.getConsumerStats().pointerBasedLag() == 0L; | 
|  | 280 | +        }); | 
|  | 281 | + | 
|  | 282 | +        // Add messages to the file | 
|  | 283 | +        try ( | 
|  | 284 | +            BufferedWriter writer = Files.newBufferedWriter( | 
|  | 285 | +                shardFile, | 
|  | 286 | +                StandardCharsets.UTF_8, | 
|  | 287 | +                StandardOpenOption.WRITE, | 
|  | 288 | +                StandardOpenOption.TRUNCATE_EXISTING | 
|  | 289 | +            ) | 
|  | 290 | +        ) { | 
|  | 291 | +            writer.write("{\"_id\":\"1\",\"_version\":\"1\",\"_op_type\":\"index\",\"_source\":{\"name\":\"alice\", \"age\": 30}}\n"); | 
|  | 292 | +            writer.write("{\"_id\":\"2\",\"_version\":\"1\",\"_op_type\":\"index\",\"_source\":{\"name\":\"bob\", \"age\": 35}}\n"); | 
|  | 293 | +            writer.flush(); | 
|  | 294 | +        } | 
|  | 295 | + | 
|  | 296 | +        try (FileChannel channel = FileChannel.open(shardFile, StandardOpenOption.READ)) { | 
|  | 297 | +            channel.force(true); | 
|  | 298 | +        } | 
|  | 299 | + | 
|  | 300 | +        // Wait for messages to be processed | 
|  | 301 | +        waitForState(() -> { | 
|  | 302 | +            SearchResponse response = client().prepareSearch(index).setQuery(new RangeQueryBuilder("age").gte(0)).get(); | 
|  | 303 | +            return response.getHits().getTotalHits().value() == 2; | 
|  | 304 | +        }); | 
|  | 305 | + | 
|  | 306 | +        // Lag should be 0 after all messages are consumed | 
|  | 307 | +        waitForState(() -> { | 
|  | 308 | +            PollingIngestStats stats = getPollingIngestStats(index); | 
|  | 309 | +            return stats != null && stats.getConsumerStats().pointerBasedLag() == 0L; | 
|  | 310 | +        }); | 
|  | 311 | + | 
|  | 312 | +        // cleanup | 
|  | 313 | +        client().admin().indices().delete(new DeleteIndexRequest(index)).actionGet(); | 
|  | 314 | +    } | 
|  | 315 | + | 
|  | 316 | +    public void testPointerBasedLagAfterPause() throws Exception { | 
|  | 317 | +        String mappings = """ | 
|  | 318 | +            { | 
|  | 319 | +              "properties": { | 
|  | 320 | +                "name": { "type": "text" }, | 
|  | 321 | +                "age": { "type": "integer" } | 
|  | 322 | +              } | 
|  | 323 | +            } | 
|  | 324 | +            """; | 
|  | 325 | + | 
|  | 326 | +        createIndexWithMappingSource( | 
|  | 327 | +            index, | 
|  | 328 | +            Settings.builder() | 
|  | 329 | +                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) | 
|  | 330 | +                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) | 
|  | 331 | +                .put("ingestion_source.type", "FILE") | 
|  | 332 | +                .put("ingestion_source.pointer.init.reset", "earliest") | 
|  | 333 | +                .put("ingestion_source.param.stream", stream) | 
|  | 334 | +                .put("ingestion_source.param.base_directory", ingestionDir.toString()) | 
|  | 335 | +                .put("index.replication.type", "SEGMENT") | 
|  | 336 | +                .build(), | 
|  | 337 | +            mappings | 
|  | 338 | +        ); | 
|  | 339 | +        ensureGreen(index); | 
|  | 340 | + | 
|  | 341 | +        // Wait for initial messages to be processed | 
|  | 342 | +        waitForState(() -> { | 
|  | 343 | +            SearchResponse response = client().prepareSearch(index).setQuery(new RangeQueryBuilder("age").gte(0)).get(); | 
|  | 344 | +            return response.getHits().getTotalHits().value() == 2; | 
|  | 345 | +        }); | 
|  | 346 | + | 
|  | 347 | +        // Pause ingestion | 
|  | 348 | +        PauseIngestionResponse pauseResponse = client().admin().indices().pauseIngestion(Requests.pauseIngestionRequest(index)).get(); | 
|  | 349 | +        assertTrue(pauseResponse.isAcknowledged()); | 
|  | 350 | +        assertTrue(pauseResponse.isShardsAcknowledged()); | 
|  | 351 | + | 
|  | 352 | +        // Wait for pause to take effect | 
|  | 353 | +        waitForState(() -> { | 
|  | 354 | +            GetIngestionStateResponse ingestionState = client().admin() | 
|  | 355 | +                .indices() | 
|  | 356 | +                .getIngestionState(Requests.getIngestionStateRequest(index)) | 
|  | 357 | +                .get(); | 
|  | 358 | +            return ingestionState.getFailedShards() == 0 | 
|  | 359 | +                && Arrays.stream(ingestionState.getShardStates()) | 
|  | 360 | +                    .allMatch(state -> state.isPollerPaused() && state.getPollerState().equalsIgnoreCase("paused")); | 
|  | 361 | +        }); | 
|  | 362 | + | 
|  | 363 | +        // Add more messages to the file while paused | 
|  | 364 | +        Path streamDir = ingestionDir.resolve(stream); | 
|  | 365 | +        Path shardFile = streamDir.resolve("0.ndjson"); | 
|  | 366 | +        try (BufferedWriter writer = Files.newBufferedWriter(shardFile, StandardCharsets.UTF_8, StandardOpenOption.APPEND)) { | 
|  | 367 | +            writer.write("{\"_id\":\"3\",\"_version\":\"1\",\"_op_type\":\"index\",\"_source\":{\"name\":\"charlie\", \"age\": 40}}\n"); | 
|  | 368 | +            writer.write("{\"_id\":\"4\",\"_version\":\"1\",\"_op_type\":\"index\",\"_source\":{\"name\":\"diana\", \"age\": 45}}\n"); | 
|  | 369 | +            writer.write("{\"_id\":\"5\",\"_version\":\"1\",\"_op_type\":\"index\",\"_source\":{\"name\":\"eve\", \"age\": 50}}\n"); | 
|  | 370 | +            writer.flush(); | 
|  | 371 | +        } | 
|  | 372 | + | 
|  | 373 | +        try (FileChannel channel = FileChannel.open(shardFile, StandardOpenOption.READ)) { | 
|  | 374 | +            channel.force(true); | 
|  | 375 | +        } | 
|  | 376 | + | 
|  | 377 | +        // Wait for lag to be calculated (lag is updated every 10 seconds) | 
|  | 378 | +        waitForState(() -> { | 
|  | 379 | +            PollingIngestStats stats = getPollingIngestStats(index); | 
|  | 380 | +            return stats != null && stats.getConsumerStats().pointerBasedLag() == 3L; | 
|  | 381 | +        }); | 
|  | 382 | + | 
|  | 383 | +        // cleanup | 
|  | 384 | +        client().admin().indices().delete(new DeleteIndexRequest(index)).actionGet(); | 
|  | 385 | +    } | 
|  | 386 | + | 
|  | 387 | +    /** | 
|  | 388 | +     * Helper method to get polling ingest stats for the index | 
|  | 389 | +     */ | 
|  | 390 | +    private PollingIngestStats getPollingIngestStats(String indexName) { | 
|  | 391 | +        IndexStats indexStats = client().admin().indices().prepareStats(indexName).get().getIndex(indexName); | 
|  | 392 | +        ShardStats[] shards = indexStats.getShards(); | 
|  | 393 | +        if (shards.length > 0) { | 
|  | 394 | +            return shards[0].getPollingIngestStats(); | 
|  | 395 | +        } | 
|  | 396 | +        return null; | 
|  | 397 | +    } | 
|  | 398 | + | 
|  | 399 | +    private void waitForState(Callable<Boolean> checkState) throws Exception { | 
|  | 400 | +        assertBusy(() -> { | 
|  | 401 | +            if (checkState.call() == false) { | 
|  | 402 | +                fail("Provided state requirements not met"); | 
|  | 403 | +            } | 
|  | 404 | +        }, 1, TimeUnit.MINUTES); | 
|  | 405 | +    } | 
| 240 | 406 | } | 
0 commit comments