From 1fe8a02fdec5e2f16eed91b2cd66234df9a09ea6 Mon Sep 17 00:00:00 2001 From: Mohit Yadav <105265192+mohityadav766@users.noreply.github.com> Date: Tue, 3 Mar 2026 16:39:27 +0530 Subject: [PATCH] Improve indexing (#26154) * Add Prometheus metrics for reindexing pipeline via Micrometer Bridge the existing reindexing atomic counters to Prometheus so operators can alert on failures, latency spikes, and backpressure without relying solely on database-flushed stats. - Add ReindexingMetrics singleton (initialize/getInstance pattern matching CacheMetrics) with job lifecycle counters, stage success/failed/warnings counters, bulk request timers with SLA buckets, payload size distribution, backpressure and promotion counters, and active/pending gauges - Register in MicrometerBundle after StreamableLogsMetrics - Instrument ReindexingOrchestrator.run() with job started/completed/failed/stopped - Bridge StageStatsTracker.flush() deltas to Prometheus per stage and entity type - Add bulk request latency timer and payload size recording in OpenSearchBulkSink - Record backpressure events in SearchIndexExecutor.handleBackpressure() - Record promotion success/failure in DefaultRecreateHandler - Add ReindexingMetricsTest with 24 tests covering all metric types * Add Improvements * Auto Gene * Use Auto Config in distributed * Fix Partition Claim Spread * Make partition use config * Correct total count * Fix Wait time to 5 mins * Revert om yaml * Fix Sink sync * Add Failure Handling at different stages * Update script to create entities * Move to scripts * Add usage and fix script * Fix Script * Update generated TypeScript types * Fix Staging miss * Fix Stats reconcilation issue * Revert workflow handler * Fix Partition worker early sync * Update Logs * Update logs EntityRepository * Error failure test * Review Comments fix * Fix Non Distributed live feed * Fix Non Distributed stats feed * Fix Review comments * Fix Time Series cutt off * Update generated TypeScript types * Md * Benchmark addition * Fix date time warning * Update load test to do benchmark analysis * Disagnostic and update perf test * Move load test to bin * Fix Review Comments * Add numeric values * Move to localhost by default * Fix Perf test issues * Review Comments * Add Preflight Fixes * Add Preflight fixes for stale entry * Remove stale entry on ApplicationHandler --------- Co-authored-by: github-actions[bot] (cherry picked from commit b59aa7fc44e9db0fc8d8d070604e13abbcf1ea6f) --- bin/distributed-test/DIAGNOSTICS.md | 274 ++ bin/distributed-test/USAGE.md | 173 ++ .../distributed-test/scripts/logs.sh | 0 bin/distributed-test/scripts/perf-test.sh | 2700 +++++++++++++++++ .../distributed-test/scripts/start.sh | 2 +- .../distributed-test/scripts/stop.sh | 0 .../scripts/trigger-reindex.sh | 0 docker/development/distributed-test/README.md | 8 +- .../scripts/load-test-data.sh | 960 ------ .../service/OpenMetadataApplication.java | 2 + .../service/apps/ApplicationHandler.java | 5 +- .../bundles/searchIndex/AdaptiveBackoff.java | 35 + .../searchIndex/BulkCircuitBreaker.java | 127 + .../apps/bundles/searchIndex/BulkSink.java | 46 +- .../CompositeProgressListener.java | 45 + .../DistributedIndexingStrategy.java | 649 ++++ .../searchIndex/ElasticSearchBulkSink.java | 87 +- .../searchIndex/EntityBatchSizeEstimator.java | 38 + .../bundles/searchIndex/EntityPriority.java | 125 + .../bundles/searchIndex/EntityReader.java | 336 ++ .../bundles/searchIndex/ExecutionResult.java | 56 +- .../searchIndex/IndexingFailureRecorder.java | 5 + .../bundles/searchIndex/IndexingPipeline.java | 534 ++++ .../bundles/searchIndex/IndexingStrategy.java | 21 + .../searchIndex/OpenSearchBulkSink.java | 145 +- .../searchIndex/OrchestratorContext.java | 32 + .../QuartzOrchestratorContext.java | 98 + .../searchIndex/ReindexingConfiguration.java | 97 +- .../searchIndex/ReindexingMetrics.java | 312 ++ .../searchIndex/ReindexingOrchestrator.java | 459 +++ .../ReindexingProgressListener.java | 23 + .../bundles/searchIndex/SearchIndexApp.java | 936 +----- .../searchIndex/SearchIndexExecutor.java | 637 ++-- .../SingleServerIndexingStrategy.java | 41 + .../bundles/searchIndex/StatsReconciler.java | 19 +- .../searchIndex/VectorCompletionResult.java | 12 + .../DistributedJobParticipant.java | 8 +- .../DistributedSearchIndexCoordinator.java | 57 +- .../DistributedSearchIndexExecutor.java | 98 +- .../distributed/EntityCompletionTracker.java | 40 +- .../distributed/PartitionCalculator.java | 100 +- .../distributed/PartitionWorker.java | 166 +- .../listeners/LoggingProgressListener.java | 32 + .../listeners/QuartzProgressListener.java | 81 +- .../searchIndex/stats/StageStatsTracker.java | 36 +- .../service/jdbi3/CollectionDAO.java | 28 +- .../service/jdbi3/EntityRepository.java | 5 + .../jdbi3/HikariCPDataSourceFactory.java | 2 +- .../service/monitoring/MicrometerBundle.java | 3 + .../resources/system/DiagnosticsResource.java | 282 ++ .../search/DefaultRecreateHandler.java | 138 +- .../service/search/IndexManagementClient.java | 19 + .../PaginatedEntityTimeSeriesSource.java | 15 + .../data/app/SearchIndexingApplication.json | 3 +- .../searchIndex/AdaptiveBackoffTest.java | 72 + .../searchIndex/BulkCircuitBreakerTest.java | 171 ++ .../EntityBatchSizeEstimatorTest.java | 67 + .../searchIndex/EntityPriorityTest.java | 151 + .../searchIndex/EntityReaderRetryTest.java | 108 + .../ReindexErrorScenarioIntegrationTest.java | 785 +++++ ...ReindexingConfigurationTimeSeriesTest.java | 137 + .../searchIndex/ReindexingMetricsTest.java | 445 +++ .../searchIndex/SearchIndexAppTest.java | 121 +- .../SearchIndexFailureIntegrationTest.java | 28 +- .../searchIndex/StatsThreadSafetyTest.java | 196 ++ ...DistributedSearchIndexCoordinatorTest.java | 17 +- ...DistributedSearchIndexIntegrationTest.java | 5 +- .../internal/searchIndexingAppConfig.json | 15 + .../json/schema/system/eventPublisherJob.json | 16 + .../Applications/SearchIndexingApplication.md | 14 + .../createIngestionPipeline.ts | 11 + .../src/generated/entity/applications/app.ts | 11 + .../internal/searchIndexingAppConfig.ts | 11 + .../marketplace/appMarketPlaceDefinition.ts | 11 + .../createAppMarketPlaceDefinitionReq.ts | 11 + .../ingestionPipelines/ingestionPipeline.ts | 11 + .../metadataIngestion/application.ts | 11 + .../metadataIngestion/applicationPipeline.ts | 11 + .../generated/metadataIngestion/workflow.ts | 11 + .../src/generated/system/eventPublisherJob.ts | 15 +- .../SearchIndexingApplication.json | 15 + 81 files changed, 10223 insertions(+), 2406 deletions(-) create mode 100644 bin/distributed-test/DIAGNOSTICS.md create mode 100644 bin/distributed-test/USAGE.md rename {docker/development => bin}/distributed-test/scripts/logs.sh (100%) create mode 100755 bin/distributed-test/scripts/perf-test.sh rename {docker/development => bin}/distributed-test/scripts/start.sh (97%) rename {docker/development => bin}/distributed-test/scripts/stop.sh (100%) rename {docker/development => bin}/distributed-test/scripts/trigger-reindex.sh (100%) delete mode 100755 docker/development/distributed-test/scripts/load-test-data.sh create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoff.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreaker.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriority.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetrics.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/VectorCompletionResult.java create mode 100644 openmetadata-service/src/main/java/org/openmetadata/service/resources/system/DiagnosticsResource.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreakerTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriorityTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexErrorScenarioIntegrationTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTimeSeriesTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetricsTest.java create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/StatsThreadSafetyTest.java diff --git a/bin/distributed-test/DIAGNOSTICS.md b/bin/distributed-test/DIAGNOSTICS.md new file mode 100644 index 00000000000..7fff96472fd --- /dev/null +++ b/bin/distributed-test/DIAGNOSTICS.md @@ -0,0 +1,274 @@ +# Server-Side Diagnostics & Load Test Correlation + +The diagnostics endpoint (`GET /api/v1/system/diagnostics`) provides a single-call performance snapshot of the OpenMetadata server. Combined with the load test script, it enables pinpointing **where** time is spent during high-load scenarios and produces actionable tuning recommendations. + +## The Diagnostics Endpoint + +### Basic Usage + +```bash +curl -H "Authorization: Bearer $TOKEN" \ + http://localhost:8585/api/v1/system/diagnostics | python3 -m json.tool +``` + +### Response Structure + +```json +{ + "timestamp": "2026-03-02T19:00:00Z", + "jvm": { ... }, + "jetty": { ... }, + "database": { ... }, + "bulk_executor": { ... }, + "request_latency": { ... } +} +``` + +Each section is explained below. + +--- + +## Understanding Each Section + +### JVM + +| Field | What It Tells You | +|-------|-------------------| +| `heap_used_bytes` / `heap_max_bytes` | Current heap consumption vs maximum. | +| `heap_usage_pct` | If >85% after load, GC pressure is likely adding tail latency. | +| `gc_pause_total_ms` | Cumulative GC pause time since JVM start. Compare before/after load to see how much GC occurred during the test. | +| `gc_count` | Total GC collections. A large delta during load means frequent stop-the-world pauses. | +| `thread_count` / `thread_peak` | Active JVM threads. Correlate with Jetty thread pool. | +| `cpu_process_pct` | Process CPU utilization (0-100). If pinned at 100%, the server is CPU-bound. | +| `uptime_seconds` | Useful to confirm the server wasn't restarted mid-test. | + +### Jetty (Thread Pool) + +| Field | What It Tells You | +|-------|-------------------| +| `threads_busy` / `threads_max` | How many request-handling threads are in use. | +| `utilization_pct` | **Key metric.** If >90% with `queue_size > 0`, the thread pool is saturated and requests are queuing. | +| `queue_size` | Requests waiting for a free thread. Non-zero means latency is being added by queuing. | +| `queue_time_avg_ms` | Average time a request waits in the queue before getting a thread. | +| `virtual_threads_enabled` | Whether Java 21 virtual threads are active (eliminates thread pool as a bottleneck). | + +### Database (HikariCP Pool) + +| Field | What It Tells You | +|-------|-------------------| +| `pool_active` / `pool_max` | Active DB connections vs maximum pool size. | +| `pool_usage_pct` | **Key metric.** If >80%, connection contention is likely. Requests wait for a free connection. | +| `pool_pending` | Threads waiting for a DB connection. If >0 during load, the pool is undersized. | +| `pool_idle` | Spare connections. If 0 during load, the pool is fully utilized. | + +### Bulk Executor + +| Field | What It Tells You | +|-------|-------------------| +| `queue_depth` / `queue_capacity` | Items in the async processing queue. | +| `queue_usage_pct` | If >70%, approaching the rejection threshold (HTTP 503 errors). | +| `active_threads` / `max_threads` | Worker threads actively processing bulk operations. | +| `has_capacity` | `false` means the next bulk submission will be rejected with 503. | + +### Request Latency (Per-Endpoint Breakdown) + +This is the most actionable section. For each `METHOD /endpoint` combination: + +| Field | What It Tells You | +|-------|-------------------| +| `count` | Total requests processed for this endpoint. | +| `avg_total_ms` | Average end-to-end latency. | +| `avg_db_ms` / `db_pct` | Time spent in database queries and its percentage of total. | +| `avg_search_ms` / `search_pct` | Time spent in search/Elasticsearch operations. | +| `avg_internal_ms` / `internal_pct` | Time in Java code (serialization, validation, business logic). | +| `avg_db_ops` / `avg_search_ops` | Average number of DB/search round-trips per request. | + +**Reading the breakdown:** If `PUT /v1/tables` shows `db_pct: 56%`, then 56% of the request time is spent waiting for database queries. Combined with `database.pool_usage_pct: 85%`, this tells you the DB connection pool is the bottleneck. + +--- + +## Load Test Integration + +The load test script automatically queries the diagnostics endpoint at three points: + +1. **Before load** — baseline snapshot +2. **During load** — sampled every 10 seconds by the health monitor +3. **After load** — final snapshot for comparison + +### Running a Load Test with Diagnostics + +```bash +# Basic: diagnostics are collected automatically +./perf-test.sh --scale small --server http://localhost:8585 --admin-port 8586 + +# With explicit token +./perf-test.sh --scale medium --server http://localhost:8585 \ + --admin-port 8586 --token "$MY_TOKEN" --output /tmp/bench.json +``` + +The `--admin-port` flag enables both Prometheus scraping and diagnostics collection. Diagnostics work without it too (they use the main API port). + +### Console Output + +After the benchmark table, you'll see a `SERVER-SIDE BREAKDOWN` section: + +``` +SERVER-SIDE BREAKDOWN (from /api/v1/system/diagnostics): + JVM: heap 1.2GB/2GB (60%), GC pauses +450ms during load + Jetty: 142/150 threads busy (95%), queue depth: 23 + DB Pool: 85/100 active (85%), 12 pending connections + Bulk Executor: queue 450/1000 (45%) + + Latency Breakdown (PUT endpoints): + Endpoint Total DB% Search% Internal% + /v1/tables 320ms 56.2% 14.1% 29.7% + /v1/topics 180ms 48.0% 22.0% 30.0% + /v1/dashboards 250ms 52.0% 18.0% 30.0% + + BOTTLENECK: DB bottleneck on PUT /v1/tables: 56.2% of request time in DB, pool at 85.0% utilization +``` + +### JSON Report + +The report includes top-level `diagnostics_before` and `diagnostics_after` objects, plus `cluster_sizing.server_side_analysis`: + +```bash +cat /tmp/bench.json | python3 -c " +import json, sys +r = json.load(sys.stdin) + +# Check if diagnostics were available +diag = r.get('diagnostics_after', {}) +if diag: + jvm = diag['jvm'] + print(f'Heap: {jvm[\"heap_usage_pct\"]}%') + print(f'GC pauses: {jvm[\"gc_pause_total_ms\"]}ms') + + jetty = diag['jetty'] + print(f'Jetty: {jetty[\"threads_busy\"]}/{jetty[\"threads_max\"]} ({jetty[\"utilization_pct\"]}%)') + + db = diag['database'] + print(f'DB pool: {db[\"pool_active\"]}/{db[\"pool_max\"]} ({db[\"pool_usage_pct\"]}%)') + + for ep, data in diag.get('request_latency', {}).items(): + print(f'{ep}: total={data[\"avg_total_ms\"]}ms ' + f'DB={data[\"db_pct\"]}% Search={data[\"search_pct\"]}% ' + f'Internal={data[\"internal_pct\"]}%') +else: + print('Diagnostics not available (server may be older version)') +" +``` + +--- + +## Bottleneck Detection Rules + +The load test applies these rules automatically and surfaces them in findings: + +| Condition | Diagnosis | Recommended Fix | +|-----------|-----------|-----------------| +| `db_pct > 60%` AND `pool_usage_pct > 80%` | DB is the bottleneck | `export DB_CONNECTION_POOL_MAX_SIZE=150` | +| `jetty.utilization_pct > 90%` AND `queue_size > 0` | Thread pool saturated | `export SERVER_MAX_THREADS=300` or enable virtual threads | +| `search_pct > 30%` for any endpoint | Search indexing consuming latency | `export ELASTICSEARCH_MAX_CONN_TOTAL=50` | +| `bulk_executor.queue_usage_pct > 70%` | Near bulk rejection threshold | `export BULK_OPERATION_QUEUE_SIZE=2000` | +| `jvm.heap_usage_pct > 85%` after load | Memory pressure / GC tail latency | Increase JVM heap (`-Xmx`) | + +--- + +## Common Scenarios + +### Scenario 1: High Latency, DB is the Bottleneck + +**Symptoms:** p95 latency >2s, `db_pct` >60%, `pool_usage_pct` >80%. + +``` + Latency Breakdown: + /v1/tables 320ms DB=62% Search=12% Internal=26% + DB Pool: 95/100 active (95%), 8 pending +``` + +**What's happening:** Every PUT requires multiple DB round-trips. At 95% pool utilization with 8 pending connections, requests are waiting for a free connection. + +**Fix:** +```bash +export DB_CONNECTION_POOL_MAX_SIZE=150 +export DB_CONNECTION_TIMEOUT=10000 # Fail fast instead of waiting 30s +``` + +### Scenario 2: Thread Pool Exhaustion + +**Symptoms:** Connection refused errors, `utilization_pct` >95%, `queue_size` growing. + +``` + Jetty: 148/150 threads busy (99%), queue depth: 45 +``` + +**What's happening:** All Jetty threads are busy. New requests queue up, adding latency. If the queue fills, connections get refused. + +**Fix:** +```bash +export SERVER_MAX_THREADS=300 +# OR enable virtual threads (preferred for I/O-bound workloads): +export SERVER_ENABLE_VIRTUAL_THREAD=true +``` + +### Scenario 3: GC Pressure + +**Symptoms:** Periodic latency spikes, `heap_usage_pct` >85%, large GC pause delta. + +``` + JVM: heap 1.8GB/2GB (90%), GC pauses +2300ms during load +``` + +**What's happening:** The JVM is spending significant time in garbage collection. This manifests as periodic latency spikes and throughput drops. + +**Fix:** +```bash +# Increase heap +export OPENMETADATA_HEAP_OPTS="-Xmx4g -Xms4g" +``` + +### Scenario 4: Bulk Executor Queue Filling + +**Symptoms:** HTTP 503 errors on PUT endpoints. + +``` + Bulk Executor: queue 980/1000 (98%) + has_capacity: false +``` + +**What's happening:** The async processing queue is full. New requests that need bulk processing are rejected with 503. + +**Fix:** +```bash +export BULK_OPERATION_QUEUE_SIZE=2000 +export BULK_OPERATION_MAX_THREADS=20 +``` + +--- + +## Comparing Before/After Snapshots + +The most valuable analysis comes from comparing diagnostics before and after load: + +| Metric | Before | After | Interpretation | +|--------|--------|-------|----------------| +| `heap_usage_pct` | 25% | 85% | Significant memory allocation during load | +| `gc_pause_total_ms` | 200 | 2500 | 2.3s of GC pauses during the test | +| `pool_active` | 2 | 95 | Pool went from idle to near-max | +| `pool_pending` | 0 | 8 | Connection contention appeared | +| `queue_depth` | 0 | 450 | Bulk queue built up under load | + +If the `diagnostics_during` samples are available in the health monitor data, you can plot these metrics over time to see exactly when bottlenecks emerged. + +--- + +## Graceful Fallback + +If the server doesn't have the diagnostics endpoint (older version), the load test: +- Prints a notice: `Diagnostics endpoint returned status=404 (may not be available)` +- Falls back to Prometheus scraping (if `--admin-port` is set) +- Skips the `SERVER-SIDE BREAKDOWN` section in the console output +- Omits `diagnostics_before`/`diagnostics_after` from the JSON report + +No hard dependency — the load test works with or without it. diff --git a/bin/distributed-test/USAGE.md b/bin/distributed-test/USAGE.md new file mode 100644 index 00000000000..02707aee62f --- /dev/null +++ b/bin/distributed-test/USAGE.md @@ -0,0 +1,173 @@ +# Distributed Indexing Load Test Scripts + +Scripts for generating test data and triggering reindexing to load-test the OpenMetadata search indexing pipeline. + +## Quick Start + +```bash +# 1. Start the environment +./scripts/start.sh + +# 2. Load test data (~50K entities) +./scripts/perf-test.sh --scale small --server http://localhost:8585 + +# 3. Trigger reindex +./scripts/trigger-reindex.sh + +# 4. Monitor logs +./scripts/logs.sh + +# 5. Stop the environment +./scripts/stop.sh +``` + +## perf-test.sh + +Generates entities across 30+ entity types, including time-series data, lineage edges, and data quality entities. Uses concurrent workers for high throughput. + +### Scale Presets + +Use `--scale` to pick a preset: + +| Preset | Approximate Total | Use Case | +|--------|-------------------|----------| +| `small` | ~50K | Quick smoke tests, CI | +| `medium` | ~500K | Integration testing | +| `large` | ~2M | Performance validation | +| `xlarge` | ~5M | Full-scale load testing | + +```bash +# Small smoke test +./perf-test.sh --scale small --server http://localhost:8585 + +# Full 5M load test +./perf-test.sh --scale xlarge --server http://localhost:8585 + +# Quick mode (~10K, fastest) +./perf-test.sh --quick --server http://localhost:8585 +``` + +Default (no `--scale` or `--quick`) produces ~46K entities for backward compatibility. + +### Overriding Individual Counts + +Any `--entity-type NUM` flag overrides the preset for that entity type: + +```bash +# Small preset but with 100K tables +./perf-test.sh --scale small --tables 100000 + +# Only create tables and dashboards (everything else stays at preset counts) +./perf-test.sh --scale small --tables 50000 --dashboards 10000 +``` + +### All Flags + +#### Entity counts + +| Flag | Default | Description | +|------|---------|-------------| +| `--tables NUM` | 20000 | Database tables | +| `--topics NUM` | 3000 | Kafka/messaging topics | +| `--dashboards NUM` | 5000 | Looker dashboards | +| `--charts NUM` | 10000 | Dashboard charts | +| `--pipelines NUM` | 3000 | Airflow pipelines | +| `--stored-procedures NUM` | 0 | Stored procedures | +| `--containers NUM` | 2000 | S3 containers | +| `--search-indexes NUM` | 1000 | Elasticsearch indexes | +| `--mlmodels NUM` | 2000 | ML models | +| `--queries NUM` | 0 | SQL queries | +| `--data-models NUM` | 0 | Dashboard data models | +| `--test-suites NUM` | 0 | Test suites | +| `--test-cases NUM` | 0 | Test cases (linked to tables) | +| `--glossaries NUM` | 50 | Glossaries | +| `--glossary-terms NUM` | 5000 | Glossary terms | +| `--classifications NUM` | 20 | Tag classifications | +| `--tags NUM` | 1000 | Tags | +| `--users NUM` | 0 | Users | +| `--teams NUM` | 0 | Teams | +| `--domains NUM` | 0 | Domains | +| `--data-products NUM` | 0 | Data products (need domains) | +| `--api-collections NUM` | 0 | API collections | +| `--api-endpoints NUM` | 0 | API endpoints (need collections) | +| `--lineage-edges NUM` | 0 | Lineage edges between entities | + +#### Time-series entity counts + +| Flag | Default | Description | +|------|---------|-------------| +| `--test-case-results NUM` | 0 | Test case results (need test cases) | +| `--entity-report-data NUM` | 0 | Entity report data insights | +| `--web-analytic-views NUM` | 0 | Web analytic entity view reports | +| `--web-analytic-activity NUM` | 0 | Web analytic user activity reports | +| `--raw-cost-analysis NUM` | 0 | Raw cost analysis reports | +| `--aggregated-cost-analysis NUM` | 0 | Aggregated cost analysis reports | + +#### Other options + +| Flag | Default | Description | +|------|---------|-------------| +| `--server URL` | `http://localhost:8585` | Target OpenMetadata server | +| `--workers NUM` | 20 | Concurrent HTTP workers | +| `--quick` | - | Quick mode preset (~10K entities) | +| `--scale PRESET` | - | Scale preset (small/medium/large/xlarge) | + +### Entity Creation Order + +The script creates entities in dependency order across 7 phases: + +``` +Phase 1 Metadata domains, classifications, tags, glossaries, terms, users, teams +Phase 2 Services database, dashboard, pipeline, messaging, ML, storage, search, API +Phase 3 Infrastructure databases, schemas, API collections +Phase 4 Core entities tables, dashboards, charts, topics, pipelines, storedProcedures, + containers, searchIndexes, mlmodels, queries, dataModels, + apiEndpoints, dataProducts +Phase 5 Data Quality testSuites, testCases +Phase 6 Lineage table->table (60%), table->dashboard (25%), pipeline->table (15%) +Phase 7 Time-Series testCaseResults, entityReportData, webAnalyticViews, + webAnalyticActivity, rawCostAnalysis, aggCostAnalysis +``` + +### Entity Linking + +- **Tables, dashboards, pipelines**: IDs collected during Phase 4 for use in lineage (Phase 6) +- **Test cases**: FQNs collected for testCaseResult creation (Phase 7) +- **Lineage edges**: Use collected UUIDs via `PUT /api/v1/lineage` +- Collections are capped at `max(lineage_edges * 2, test_case_results)` to bound memory + +### Auto-Scaling Infrastructure + +Databases and schemas scale automatically with table count: +- `NUM_DATABASES = max(1, tables / 50000)` +- `SCHEMAS_PER_DB = min(20, tables / (databases * 5000))` +- This keeps ~5000 tables per schema at any scale + +### Retry Logic + +HTTP requests retry up to 3 times with exponential backoff (1s, 2s, 4s) on: +- 5xx server errors +- Connection errors / timeouts + +### Performance Tips + +- Use `--workers 30` or higher if the server can handle it +- Time-series and lineage phases use `min(10, workers)` to avoid overwhelming the server +- At `xlarge` scale, expect the script to run for several hours depending on server capacity +- Monitor server logs for 429/503 errors and reduce workers if needed + +## Verification After Loading + +```bash +# 1. Trigger reindex +./scripts/trigger-reindex.sh + +# 2. Check partition table for all entity types +mysql -e "SELECT DISTINCT entityType FROM search_index_partition ORDER BY entityType;" + +# 3. Verify counts in UI +# - Data Assets: tables, topics, dashboards, pipelines, etc. +# - Data Quality: test suites and test cases +# - Lineage: visible edges between tables/dashboards/pipelines +# - Data Insights: time-series charts for entity reports, web analytics, cost analysis +``` diff --git a/docker/development/distributed-test/scripts/logs.sh b/bin/distributed-test/scripts/logs.sh similarity index 100% rename from docker/development/distributed-test/scripts/logs.sh rename to bin/distributed-test/scripts/logs.sh diff --git a/bin/distributed-test/scripts/perf-test.sh b/bin/distributed-test/scripts/perf-test.sh new file mode 100755 index 00000000000..e6a9bbf25f7 --- /dev/null +++ b/bin/distributed-test/scripts/perf-test.sh @@ -0,0 +1,2700 @@ +#!/bin/bash +# Load test data for distributed indexing testing +# Supports 30+ entity types including time-series data, lineage, and data quality +# Includes benchmarking, latency tracking, and cluster sizing recommendations + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# ── Default values (backward-compatible ~46k) ──────────────────────────────── +SERVER_URL="http://localhost:8585" +NUM_TABLES=20000 +NUM_TOPICS=3000 +NUM_DASHBOARDS=5000 +NUM_CHARTS=10000 +NUM_PIPELINES=3000 +NUM_STORED_PROCEDURES=0 +NUM_CONTAINERS=2000 +NUM_SEARCH_INDEXES=1000 +NUM_MLMODELS=2000 +NUM_QUERIES=0 +NUM_DASHBOARD_DATA_MODELS=0 +NUM_TEST_SUITES=0 +NUM_TEST_CASES=0 +NUM_GLOSSARY_TERMS=5000 +NUM_GLOSSARIES=50 +NUM_TAGS=1000 +NUM_CLASSIFICATIONS=20 +NUM_USERS=0 +NUM_TEAMS=0 +NUM_DOMAINS=0 +NUM_DATA_PRODUCTS=0 +NUM_API_COLLECTIONS=0 +NUM_API_ENDPOINTS=0 +NUM_LINEAGE_EDGES=0 +NUM_TEST_CASE_RESULTS=0 +NUM_ENTITY_REPORT_DATA=0 +NUM_WEB_ANALYTIC_VIEWS=0 +NUM_WEB_ANALYTIC_ACTIVITY=0 +NUM_RAW_COST_ANALYSIS=0 +NUM_AGG_COST_ANALYSIS=0 +NUM_WORKERS=20 +SCALE_APPLIED="" +ONLY_ENTITIES="" +OUTPUT_PATH="" +AUTH_TOKEN="" +RAMP_MODE="" +RAMP_BATCH=100 +ADMIN_PORT="" + +# Parse arguments +while [[ $# -gt 0 ]]; do + case $1 in + --scale) + SCALE_APPLIED="$2" + case $2 in + xlarge) + NUM_TABLES=2500000; NUM_TOPICS=400000; NUM_DASHBOARDS=200000; NUM_CHARTS=400000 + NUM_PIPELINES=100000; NUM_STORED_PROCEDURES=100000; NUM_CONTAINERS=75000 + NUM_SEARCH_INDEXES=50000; NUM_MLMODELS=50000; NUM_QUERIES=100000 + NUM_DASHBOARD_DATA_MODELS=50000; NUM_TEST_SUITES=1500; NUM_TEST_CASES=150000 + NUM_GLOSSARY_TERMS=50000; NUM_GLOSSARIES=500; NUM_TAGS=10000; NUM_CLASSIFICATIONS=200 + NUM_USERS=5000; NUM_TEAMS=500; NUM_DOMAINS=50; NUM_DATA_PRODUCTS=500 + NUM_API_COLLECTIONS=500; NUM_API_ENDPOINTS=50000; NUM_LINEAGE_EDGES=200000 + NUM_TEST_CASE_RESULTS=300000; NUM_ENTITY_REPORT_DATA=50000 + NUM_WEB_ANALYTIC_VIEWS=100000; NUM_WEB_ANALYTIC_ACTIVITY=50000 + NUM_RAW_COST_ANALYSIS=25000; NUM_AGG_COST_ANALYSIS=25000 + ;; + large) + NUM_TABLES=1000000; NUM_TOPICS=160000; NUM_DASHBOARDS=80000; NUM_CHARTS=160000 + NUM_PIPELINES=40000; NUM_STORED_PROCEDURES=40000; NUM_CONTAINERS=30000 + NUM_SEARCH_INDEXES=20000; NUM_MLMODELS=20000; NUM_QUERIES=40000 + NUM_DASHBOARD_DATA_MODELS=20000; NUM_TEST_SUITES=600; NUM_TEST_CASES=60000 + NUM_GLOSSARY_TERMS=20000; NUM_GLOSSARIES=200; NUM_TAGS=4000; NUM_CLASSIFICATIONS=80 + NUM_USERS=2000; NUM_TEAMS=200; NUM_DOMAINS=20; NUM_DATA_PRODUCTS=200 + NUM_API_COLLECTIONS=200; NUM_API_ENDPOINTS=20000; NUM_LINEAGE_EDGES=80000 + NUM_TEST_CASE_RESULTS=120000; NUM_ENTITY_REPORT_DATA=20000 + NUM_WEB_ANALYTIC_VIEWS=40000; NUM_WEB_ANALYTIC_ACTIVITY=20000 + NUM_RAW_COST_ANALYSIS=10000; NUM_AGG_COST_ANALYSIS=10000 + ;; + medium) + NUM_TABLES=250000; NUM_TOPICS=40000; NUM_DASHBOARDS=20000; NUM_CHARTS=40000 + NUM_PIPELINES=10000; NUM_STORED_PROCEDURES=10000; NUM_CONTAINERS=7500 + NUM_SEARCH_INDEXES=5000; NUM_MLMODELS=5000; NUM_QUERIES=10000 + NUM_DASHBOARD_DATA_MODELS=5000; NUM_TEST_SUITES=150; NUM_TEST_CASES=15000 + NUM_GLOSSARY_TERMS=5000; NUM_GLOSSARIES=50; NUM_TAGS=1000; NUM_CLASSIFICATIONS=20 + NUM_USERS=500; NUM_TEAMS=50; NUM_DOMAINS=5; NUM_DATA_PRODUCTS=50 + NUM_API_COLLECTIONS=50; NUM_API_ENDPOINTS=5000; NUM_LINEAGE_EDGES=20000 + NUM_TEST_CASE_RESULTS=30000; NUM_ENTITY_REPORT_DATA=5000 + NUM_WEB_ANALYTIC_VIEWS=10000; NUM_WEB_ANALYTIC_ACTIVITY=5000 + NUM_RAW_COST_ANALYSIS=2500; NUM_AGG_COST_ANALYSIS=2500 + ;; + 10k) + NUM_TABLES=5000; NUM_TOPICS=800; NUM_DASHBOARDS=400; NUM_CHARTS=800 + NUM_PIPELINES=200; NUM_STORED_PROCEDURES=200; NUM_CONTAINERS=150 + NUM_SEARCH_INDEXES=100; NUM_MLMODELS=100; NUM_QUERIES=200 + NUM_DASHBOARD_DATA_MODELS=100; NUM_TEST_SUITES=3; NUM_TEST_CASES=300 + NUM_GLOSSARY_TERMS=100; NUM_GLOSSARIES=1; NUM_TAGS=20; NUM_CLASSIFICATIONS=1 + NUM_USERS=10; NUM_TEAMS=1; NUM_DOMAINS=1; NUM_DATA_PRODUCTS=1 + NUM_API_COLLECTIONS=1; NUM_API_ENDPOINTS=100; NUM_LINEAGE_EDGES=400 + NUM_TEST_CASE_RESULTS=600; NUM_ENTITY_REPORT_DATA=100 + NUM_WEB_ANALYTIC_VIEWS=200; NUM_WEB_ANALYTIC_ACTIVITY=100 + NUM_RAW_COST_ANALYSIS=50; NUM_AGG_COST_ANALYSIS=50 + ;; + small|50k) + NUM_TABLES=25000; NUM_TOPICS=4000; NUM_DASHBOARDS=2000; NUM_CHARTS=4000 + NUM_PIPELINES=1000; NUM_STORED_PROCEDURES=1000; NUM_CONTAINERS=750 + NUM_SEARCH_INDEXES=500; NUM_MLMODELS=500; NUM_QUERIES=1000 + NUM_DASHBOARD_DATA_MODELS=500; NUM_TEST_SUITES=15; NUM_TEST_CASES=1500 + NUM_GLOSSARY_TERMS=500; NUM_GLOSSARIES=5; NUM_TAGS=100; NUM_CLASSIFICATIONS=2 + NUM_USERS=50; NUM_TEAMS=5; NUM_DOMAINS=1; NUM_DATA_PRODUCTS=5 + NUM_API_COLLECTIONS=5; NUM_API_ENDPOINTS=500; NUM_LINEAGE_EDGES=2000 + NUM_TEST_CASE_RESULTS=3000; NUM_ENTITY_REPORT_DATA=500 + NUM_WEB_ANALYTIC_VIEWS=1000; NUM_WEB_ANALYTIC_ACTIVITY=500 + NUM_RAW_COST_ANALYSIS=250; NUM_AGG_COST_ANALYSIS=250 + ;; + 100k) + NUM_TABLES=50000; NUM_TOPICS=8000; NUM_DASHBOARDS=4000; NUM_CHARTS=8000 + NUM_PIPELINES=2000; NUM_STORED_PROCEDURES=2000; NUM_CONTAINERS=1500 + NUM_SEARCH_INDEXES=1000; NUM_MLMODELS=1000; NUM_QUERIES=2000 + NUM_DASHBOARD_DATA_MODELS=1000; NUM_TEST_SUITES=30; NUM_TEST_CASES=3000 + NUM_GLOSSARY_TERMS=1000; NUM_GLOSSARIES=10; NUM_TAGS=200; NUM_CLASSIFICATIONS=4 + NUM_USERS=100; NUM_TEAMS=10; NUM_DOMAINS=2; NUM_DATA_PRODUCTS=10 + NUM_API_COLLECTIONS=10; NUM_API_ENDPOINTS=1000; NUM_LINEAGE_EDGES=4000 + NUM_TEST_CASE_RESULTS=6000; NUM_ENTITY_REPORT_DATA=1000 + NUM_WEB_ANALYTIC_VIEWS=2000; NUM_WEB_ANALYTIC_ACTIVITY=1000 + NUM_RAW_COST_ANALYSIS=500; NUM_AGG_COST_ANALYSIS=500 + ;; + 150k) + NUM_TABLES=75000; NUM_TOPICS=12000; NUM_DASHBOARDS=6000; NUM_CHARTS=12000 + NUM_PIPELINES=3000; NUM_STORED_PROCEDURES=3000; NUM_CONTAINERS=2250 + NUM_SEARCH_INDEXES=1500; NUM_MLMODELS=1500; NUM_QUERIES=3000 + NUM_DASHBOARD_DATA_MODELS=1500; NUM_TEST_SUITES=45; NUM_TEST_CASES=4500 + NUM_GLOSSARY_TERMS=1500; NUM_GLOSSARIES=15; NUM_TAGS=300; NUM_CLASSIFICATIONS=6 + NUM_USERS=150; NUM_TEAMS=15; NUM_DOMAINS=3; NUM_DATA_PRODUCTS=15 + NUM_API_COLLECTIONS=15; NUM_API_ENDPOINTS=1500; NUM_LINEAGE_EDGES=6000 + NUM_TEST_CASE_RESULTS=9000; NUM_ENTITY_REPORT_DATA=1500 + NUM_WEB_ANALYTIC_VIEWS=3000; NUM_WEB_ANALYTIC_ACTIVITY=1500 + NUM_RAW_COST_ANALYSIS=750; NUM_AGG_COST_ANALYSIS=750 + ;; + 200k) + NUM_TABLES=100000; NUM_TOPICS=16000; NUM_DASHBOARDS=8000; NUM_CHARTS=16000 + NUM_PIPELINES=4000; NUM_STORED_PROCEDURES=4000; NUM_CONTAINERS=3000 + NUM_SEARCH_INDEXES=2000; NUM_MLMODELS=2000; NUM_QUERIES=4000 + NUM_DASHBOARD_DATA_MODELS=2000; NUM_TEST_SUITES=60; NUM_TEST_CASES=6000 + NUM_GLOSSARY_TERMS=2000; NUM_GLOSSARIES=20; NUM_TAGS=400; NUM_CLASSIFICATIONS=8 + NUM_USERS=200; NUM_TEAMS=20; NUM_DOMAINS=4; NUM_DATA_PRODUCTS=20 + NUM_API_COLLECTIONS=20; NUM_API_ENDPOINTS=2000; NUM_LINEAGE_EDGES=8000 + NUM_TEST_CASE_RESULTS=12000; NUM_ENTITY_REPORT_DATA=2000 + NUM_WEB_ANALYTIC_VIEWS=4000; NUM_WEB_ANALYTIC_ACTIVITY=2000 + NUM_RAW_COST_ANALYSIS=1000; NUM_AGG_COST_ANALYSIS=1000 + ;; + 250k) + NUM_TABLES=125000; NUM_TOPICS=20000; NUM_DASHBOARDS=10000; NUM_CHARTS=20000 + NUM_PIPELINES=5000; NUM_STORED_PROCEDURES=5000; NUM_CONTAINERS=3750 + NUM_SEARCH_INDEXES=2500; NUM_MLMODELS=2500; NUM_QUERIES=5000 + NUM_DASHBOARD_DATA_MODELS=2500; NUM_TEST_SUITES=75; NUM_TEST_CASES=7500 + NUM_GLOSSARY_TERMS=2500; NUM_GLOSSARIES=25; NUM_TAGS=500; NUM_CLASSIFICATIONS=10 + NUM_USERS=250; NUM_TEAMS=25; NUM_DOMAINS=5; NUM_DATA_PRODUCTS=25 + NUM_API_COLLECTIONS=25; NUM_API_ENDPOINTS=2500; NUM_LINEAGE_EDGES=10000 + NUM_TEST_CASE_RESULTS=15000; NUM_ENTITY_REPORT_DATA=2500 + NUM_WEB_ANALYTIC_VIEWS=5000; NUM_WEB_ANALYTIC_ACTIVITY=2500 + NUM_RAW_COST_ANALYSIS=1250; NUM_AGG_COST_ANALYSIS=1250 + ;; + 300k) + NUM_TABLES=150000; NUM_TOPICS=24000; NUM_DASHBOARDS=12000; NUM_CHARTS=24000 + NUM_PIPELINES=6000; NUM_STORED_PROCEDURES=6000; NUM_CONTAINERS=4500 + NUM_SEARCH_INDEXES=3000; NUM_MLMODELS=3000; NUM_QUERIES=6000 + NUM_DASHBOARD_DATA_MODELS=3000; NUM_TEST_SUITES=90; NUM_TEST_CASES=9000 + NUM_GLOSSARY_TERMS=3000; NUM_GLOSSARIES=30; NUM_TAGS=600; NUM_CLASSIFICATIONS=12 + NUM_USERS=300; NUM_TEAMS=30; NUM_DOMAINS=6; NUM_DATA_PRODUCTS=30 + NUM_API_COLLECTIONS=30; NUM_API_ENDPOINTS=3000; NUM_LINEAGE_EDGES=12000 + NUM_TEST_CASE_RESULTS=18000; NUM_ENTITY_REPORT_DATA=3000 + NUM_WEB_ANALYTIC_VIEWS=6000; NUM_WEB_ANALYTIC_ACTIVITY=3000 + NUM_RAW_COST_ANALYSIS=1500; NUM_AGG_COST_ANALYSIS=1500 + ;; + 350k) + NUM_TABLES=175000; NUM_TOPICS=28000; NUM_DASHBOARDS=14000; NUM_CHARTS=28000 + NUM_PIPELINES=7000; NUM_STORED_PROCEDURES=7000; NUM_CONTAINERS=5250 + NUM_SEARCH_INDEXES=3500; NUM_MLMODELS=3500; NUM_QUERIES=7000 + NUM_DASHBOARD_DATA_MODELS=3500; NUM_TEST_SUITES=105; NUM_TEST_CASES=10500 + NUM_GLOSSARY_TERMS=3500; NUM_GLOSSARIES=35; NUM_TAGS=700; NUM_CLASSIFICATIONS=14 + NUM_USERS=350; NUM_TEAMS=35; NUM_DOMAINS=7; NUM_DATA_PRODUCTS=35 + NUM_API_COLLECTIONS=35; NUM_API_ENDPOINTS=3500; NUM_LINEAGE_EDGES=14000 + NUM_TEST_CASE_RESULTS=21000; NUM_ENTITY_REPORT_DATA=3500 + NUM_WEB_ANALYTIC_VIEWS=7000; NUM_WEB_ANALYTIC_ACTIVITY=3500 + NUM_RAW_COST_ANALYSIS=1750; NUM_AGG_COST_ANALYSIS=1750 + ;; + 400k) + NUM_TABLES=200000; NUM_TOPICS=32000; NUM_DASHBOARDS=16000; NUM_CHARTS=32000 + NUM_PIPELINES=8000; NUM_STORED_PROCEDURES=8000; NUM_CONTAINERS=6000 + NUM_SEARCH_INDEXES=4000; NUM_MLMODELS=4000; NUM_QUERIES=8000 + NUM_DASHBOARD_DATA_MODELS=4000; NUM_TEST_SUITES=120; NUM_TEST_CASES=12000 + NUM_GLOSSARY_TERMS=4000; NUM_GLOSSARIES=40; NUM_TAGS=800; NUM_CLASSIFICATIONS=16 + NUM_USERS=400; NUM_TEAMS=40; NUM_DOMAINS=8; NUM_DATA_PRODUCTS=40 + NUM_API_COLLECTIONS=40; NUM_API_ENDPOINTS=4000; NUM_LINEAGE_EDGES=16000 + NUM_TEST_CASE_RESULTS=24000; NUM_ENTITY_REPORT_DATA=4000 + NUM_WEB_ANALYTIC_VIEWS=8000; NUM_WEB_ANALYTIC_ACTIVITY=4000 + NUM_RAW_COST_ANALYSIS=2000; NUM_AGG_COST_ANALYSIS=2000 + ;; + 450k) + NUM_TABLES=225000; NUM_TOPICS=36000; NUM_DASHBOARDS=18000; NUM_CHARTS=36000 + NUM_PIPELINES=9000; NUM_STORED_PROCEDURES=9000; NUM_CONTAINERS=6750 + NUM_SEARCH_INDEXES=4500; NUM_MLMODELS=4500; NUM_QUERIES=9000 + NUM_DASHBOARD_DATA_MODELS=4500; NUM_TEST_SUITES=135; NUM_TEST_CASES=13500 + NUM_GLOSSARY_TERMS=4500; NUM_GLOSSARIES=45; NUM_TAGS=900; NUM_CLASSIFICATIONS=18 + NUM_USERS=450; NUM_TEAMS=45; NUM_DOMAINS=9; NUM_DATA_PRODUCTS=45 + NUM_API_COLLECTIONS=45; NUM_API_ENDPOINTS=4500; NUM_LINEAGE_EDGES=18000 + NUM_TEST_CASE_RESULTS=27000; NUM_ENTITY_REPORT_DATA=4500 + NUM_WEB_ANALYTIC_VIEWS=9000; NUM_WEB_ANALYTIC_ACTIVITY=4500 + NUM_RAW_COST_ANALYSIS=2250; NUM_AGG_COST_ANALYSIS=2250 + ;; + 500k) + NUM_TABLES=250000; NUM_TOPICS=40000; NUM_DASHBOARDS=20000; NUM_CHARTS=40000 + NUM_PIPELINES=10000; NUM_STORED_PROCEDURES=10000; NUM_CONTAINERS=7500 + NUM_SEARCH_INDEXES=5000; NUM_MLMODELS=5000; NUM_QUERIES=10000 + NUM_DASHBOARD_DATA_MODELS=5000; NUM_TEST_SUITES=150; NUM_TEST_CASES=15000 + NUM_GLOSSARY_TERMS=5000; NUM_GLOSSARIES=50; NUM_TAGS=1000; NUM_CLASSIFICATIONS=20 + NUM_USERS=500; NUM_TEAMS=50; NUM_DOMAINS=10; NUM_DATA_PRODUCTS=50 + NUM_API_COLLECTIONS=50; NUM_API_ENDPOINTS=5000; NUM_LINEAGE_EDGES=20000 + NUM_TEST_CASE_RESULTS=30000; NUM_ENTITY_REPORT_DATA=5000 + NUM_WEB_ANALYTIC_VIEWS=10000; NUM_WEB_ANALYTIC_ACTIVITY=5000 + NUM_RAW_COST_ANALYSIS=2500; NUM_AGG_COST_ANALYSIS=2500 + ;; + *) + echo "Unknown scale: $2 (use 10k|50k|100k|150k|200k|250k|300k|350k|400k|450k|500k|small|medium|large|xlarge)" + exit 1 + ;; + esac + shift 2 + ;; + --tables) NUM_TABLES="$2"; shift 2 ;; + --dashboards) NUM_DASHBOARDS="$2"; shift 2 ;; + --charts) NUM_CHARTS="$2"; shift 2 ;; + --pipelines) NUM_PIPELINES="$2"; shift 2 ;; + --topics) NUM_TOPICS="$2"; shift 2 ;; + --mlmodels) NUM_MLMODELS="$2"; shift 2 ;; + --containers) NUM_CONTAINERS="$2"; shift 2 ;; + --search-indexes) NUM_SEARCH_INDEXES="$2"; shift 2 ;; + --stored-procedures) NUM_STORED_PROCEDURES="$2"; shift 2 ;; + --queries) NUM_QUERIES="$2"; shift 2 ;; + --data-models) NUM_DASHBOARD_DATA_MODELS="$2"; shift 2 ;; + --test-suites) NUM_TEST_SUITES="$2"; shift 2 ;; + --test-cases) NUM_TEST_CASES="$2"; shift 2 ;; + --glossaries) NUM_GLOSSARIES="$2"; shift 2 ;; + --glossary-terms) NUM_GLOSSARY_TERMS="$2"; shift 2 ;; + --classifications) NUM_CLASSIFICATIONS="$2"; shift 2 ;; + --tags) NUM_TAGS="$2"; shift 2 ;; + --users) NUM_USERS="$2"; shift 2 ;; + --teams) NUM_TEAMS="$2"; shift 2 ;; + --domains) NUM_DOMAINS="$2"; shift 2 ;; + --data-products) NUM_DATA_PRODUCTS="$2"; shift 2 ;; + --api-collections) NUM_API_COLLECTIONS="$2"; shift 2 ;; + --api-endpoints) NUM_API_ENDPOINTS="$2"; shift 2 ;; + --lineage-edges) NUM_LINEAGE_EDGES="$2"; shift 2 ;; + --test-case-results) NUM_TEST_CASE_RESULTS="$2"; shift 2 ;; + --entity-report-data) NUM_ENTITY_REPORT_DATA="$2"; shift 2 ;; + --web-analytic-views) NUM_WEB_ANALYTIC_VIEWS="$2"; shift 2 ;; + --web-analytic-activity) NUM_WEB_ANALYTIC_ACTIVITY="$2"; shift 2 ;; + --raw-cost-analysis) NUM_RAW_COST_ANALYSIS="$2"; shift 2 ;; + --aggregated-cost-analysis) NUM_AGG_COST_ANALYSIS="$2"; shift 2 ;; + --workers) NUM_WORKERS="$2"; shift 2 ;; + --only) ONLY_ENTITIES="$2"; shift 2 ;; + --output) OUTPUT_PATH="$2"; shift 2 ;; + --token) AUTH_TOKEN="$2"; shift 2 ;; + --ramp) RAMP_MODE="true"; shift ;; + --ramp-batch) RAMP_BATCH="$2"; shift 2 ;; + --admin-port) ADMIN_PORT="$2"; shift 2 ;; + --databases) shift 2 ;; # ignored, auto-calculated now + --terms-per-glossary) shift 2 ;; # ignored, use --glossary-terms + --tags-per-classification) shift 2 ;; # ignored, use --tags + --server) SERVER_URL="$2"; shift 2 ;; + --quick) + SCALE_APPLIED="quick" + NUM_TABLES=3000; NUM_DASHBOARDS=1000; NUM_CHARTS=2000; NUM_PIPELINES=500 + NUM_TOPICS=500; NUM_MLMODELS=300; NUM_CONTAINERS=300; NUM_SEARCH_INDEXES=200 + NUM_STORED_PROCEDURES=200; NUM_QUERIES=200; NUM_DASHBOARD_DATA_MODELS=100 + NUM_TEST_SUITES=5; NUM_TEST_CASES=500; NUM_GLOSSARY_TERMS=500; NUM_GLOSSARIES=10 + NUM_TAGS=100; NUM_CLASSIFICATIONS=5; NUM_USERS=20; NUM_TEAMS=3; NUM_DOMAINS=1 + NUM_DATA_PRODUCTS=3; NUM_API_COLLECTIONS=3; NUM_API_ENDPOINTS=100 + NUM_LINEAGE_EDGES=500; NUM_TEST_CASE_RESULTS=1000; NUM_ENTITY_REPORT_DATA=100 + NUM_WEB_ANALYTIC_VIEWS=200; NUM_WEB_ANALYTIC_ACTIVITY=100 + NUM_RAW_COST_ANALYSIS=50; NUM_AGG_COST_ANALYSIS=50 + shift + ;; + -h|--help) + cat << 'HELPEOF' +Usage: perf-test.sh [OPTIONS] + +Scale presets: + --scale {10k|50k|100k|...|500k|small|medium|large|xlarge} + Apply a preset (see table below) + --quick Quick mode (~10k entities) + +Individual entity counts override any preset: + --tables NUM --dashboards NUM --charts NUM + --pipelines NUM --topics NUM --mlmodels NUM + --containers NUM --search-indexes NUM --stored-procedures NUM + --queries NUM --data-models NUM --test-suites NUM + --test-cases NUM --glossaries NUM --glossary-terms NUM + --classifications NUM --tags NUM --users NUM + --teams NUM --domains NUM --data-products NUM + --api-collections NUM --api-endpoints NUM --lineage-edges NUM + +Time-series entity counts: + --test-case-results NUM --entity-report-data NUM + --web-analytic-views NUM --web-analytic-activity NUM + --raw-cost-analysis NUM --aggregated-cost-analysis NUM + +Benchmarking & filtering: + --only ENTITIES Comma-separated entity types to create (e.g. tables,charts,topics) + When specified, only listed entities run. Prerequisites auto-enabled. + Valid names: tables, topics, dashboards, charts, pipelines, + storedProcedures, containers, searchIndexes, mlmodels, queries, + dashboardDataModels, testSuites, testCases, glossaries, + glossaryTerms, users, teams, domains, dataProducts, + apiCollections, apiEndpoints, lineageEdges, testCaseResults, + entityReportData, webAnalyticViews, webAnalyticActivity, + rawCostAnalysis, aggCostAnalysis + --output PATH Write JSON benchmark report to PATH + (default: benchmark-report-{timestamp}.json in current dir) + --token TOKEN Auth token (overrides hardcoded default) + --ramp Run concurrency ramp test before main load + --ramp-batch NUM Entities per ramp level (default: 100) + --admin-port PORT Admin port for Prometheus metrics scraping + +Other: + --server URL Target server URL (default: https://mohitcorp.getcollate.io) + --workers NUM Concurrent workers (default: 20) + -h, --help Show this help message + +Scale preset totals (numeric): + 10k ~10K 50k ~50K 100k ~100K 150k ~150K 200k ~200K + 250k ~250K 300k ~300K 350k ~350K 400k ~400K 450k ~450K + 500k ~500K + +Scale preset totals (named): + xlarge ~5M large ~2M medium ~500K small ~50K + +Examples: + # Quick benchmark with just tables + ./perf-test.sh --only tables --tables 100 --workers 5 + + # Full benchmark with JSON output + ./perf-test.sh --quick --workers 10 --output /tmp/bench.json + + # Only tables and charts, custom counts + ./perf-test.sh --only tables,charts --tables 5000 --charts 2000 + + # Ramp test to find optimal concurrency + ./perf-test.sh --ramp --only tables --tables 500 --workers 32 + + # Full benchmark with Prometheus metrics from admin port + ./perf-test.sh --quick --workers 10 --admin-port 8586 --output /tmp/bench.json +HELPEOF + exit 0 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# Calculate total +TOTAL=$((NUM_TABLES + NUM_TOPICS + NUM_DASHBOARDS + NUM_CHARTS + NUM_PIPELINES + \ + NUM_STORED_PROCEDURES + NUM_CONTAINERS + NUM_SEARCH_INDEXES + NUM_MLMODELS + \ + NUM_QUERIES + NUM_DASHBOARD_DATA_MODELS + NUM_TEST_SUITES + NUM_TEST_CASES + \ + NUM_GLOSSARY_TERMS + NUM_GLOSSARIES + NUM_TAGS + NUM_CLASSIFICATIONS + \ + NUM_USERS + NUM_TEAMS + NUM_DOMAINS + NUM_DATA_PRODUCTS + \ + NUM_API_COLLECTIONS + NUM_API_ENDPOINTS + NUM_LINEAGE_EDGES + \ + NUM_TEST_CASE_RESULTS + NUM_ENTITY_REPORT_DATA + \ + NUM_WEB_ANALYTIC_VIEWS + NUM_WEB_ANALYTIC_ACTIVITY + \ + NUM_RAW_COST_ANALYSIS + NUM_AGG_COST_ANALYSIS)) + +echo "======================================" +echo "Loading Test Data for Distributed Indexing" +echo "======================================" +echo "Server: $SERVER_URL" +if [ -n "$SCALE_APPLIED" ]; then + echo "Scale: $SCALE_APPLIED" +fi +echo "Workers: $NUM_WORKERS" +if [ -n "$ONLY_ENTITIES" ]; then + echo "Only: $ONLY_ENTITIES" +fi +if [ -n "$RAMP_MODE" ]; then + echo "Ramp: enabled (batch=$RAMP_BATCH)" +fi +if [ -n "$ADMIN_PORT" ]; then + echo "Admin: port $ADMIN_PORT (Prometheus metrics)" +fi +echo "" +echo "Entity counts:" +printf " %-26s %s\n" "Tables:" "$NUM_TABLES" +printf " %-26s %s\n" "Topics:" "$NUM_TOPICS" +printf " %-26s %s\n" "Dashboards:" "$NUM_DASHBOARDS" +printf " %-26s %s\n" "Charts:" "$NUM_CHARTS" +printf " %-26s %s\n" "Pipelines:" "$NUM_PIPELINES" +printf " %-26s %s\n" "Stored Procedures:" "$NUM_STORED_PROCEDURES" +printf " %-26s %s\n" "Containers:" "$NUM_CONTAINERS" +printf " %-26s %s\n" "Search Indexes:" "$NUM_SEARCH_INDEXES" +printf " %-26s %s\n" "ML Models:" "$NUM_MLMODELS" +printf " %-26s %s\n" "Queries:" "$NUM_QUERIES" +printf " %-26s %s\n" "Dashboard Data Models:" "$NUM_DASHBOARD_DATA_MODELS" +printf " %-26s %s\n" "Test Suites:" "$NUM_TEST_SUITES" +printf " %-26s %s\n" "Test Cases:" "$NUM_TEST_CASES" +printf " %-26s %s\n" "Glossaries:" "$NUM_GLOSSARIES" +printf " %-26s %s\n" "Glossary Terms:" "$NUM_GLOSSARY_TERMS" +printf " %-26s %s\n" "Classifications:" "$NUM_CLASSIFICATIONS" +printf " %-26s %s\n" "Tags:" "$NUM_TAGS" +printf " %-26s %s\n" "Users:" "$NUM_USERS" +printf " %-26s %s\n" "Teams:" "$NUM_TEAMS" +printf " %-26s %s\n" "Domains:" "$NUM_DOMAINS" +printf " %-26s %s\n" "Data Products:" "$NUM_DATA_PRODUCTS" +printf " %-26s %s\n" "API Collections:" "$NUM_API_COLLECTIONS" +printf " %-26s %s\n" "API Endpoints:" "$NUM_API_ENDPOINTS" +printf " %-26s %s\n" "Lineage Edges:" "$NUM_LINEAGE_EDGES" +printf " %-26s %s\n" "Test Case Results (TS):" "$NUM_TEST_CASE_RESULTS" +printf " %-26s %s\n" "Entity Report Data (TS):" "$NUM_ENTITY_REPORT_DATA" +printf " %-26s %s\n" "Web Analytic Views (TS):" "$NUM_WEB_ANALYTIC_VIEWS" +printf " %-26s %s\n" "Web Analytic Activity (TS):" "$NUM_WEB_ANALYTIC_ACTIVITY" +printf " %-26s %s\n" "Raw Cost Analysis (TS):" "$NUM_RAW_COST_ANALYSIS" +printf " %-26s %s\n" "Agg Cost Analysis (TS):" "$NUM_AGG_COST_ANALYSIS" +echo " --------------------------" +printf " %-26s %s\n" "Total:" "$TOTAL" +echo "" + +python3 << PYEOF +import urllib.request +import urllib.error +import json +import sys +import time +import random +import uuid +import threading +import os +from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime, timedelta, timezone + +SERVER_URL = "${SERVER_URL}" +NUM_WORKERS = ${NUM_WORKERS} + +# ── Entity counts ──────────────────────────────────────────────────────────── +NUM_TABLES = ${NUM_TABLES} +NUM_TOPICS = ${NUM_TOPICS} +NUM_DASHBOARDS = ${NUM_DASHBOARDS} +NUM_CHARTS = ${NUM_CHARTS} +NUM_PIPELINES = ${NUM_PIPELINES} +NUM_STORED_PROCEDURES = ${NUM_STORED_PROCEDURES} +NUM_CONTAINERS = ${NUM_CONTAINERS} +NUM_SEARCH_INDEXES = ${NUM_SEARCH_INDEXES} +NUM_MLMODELS = ${NUM_MLMODELS} +NUM_QUERIES = ${NUM_QUERIES} +NUM_DASHBOARD_DATA_MODELS = ${NUM_DASHBOARD_DATA_MODELS} +NUM_TEST_SUITES = ${NUM_TEST_SUITES} +NUM_TEST_CASES = ${NUM_TEST_CASES} +NUM_GLOSSARY_TERMS = ${NUM_GLOSSARY_TERMS} +NUM_GLOSSARIES = ${NUM_GLOSSARIES} +NUM_TAGS = ${NUM_TAGS} +NUM_CLASSIFICATIONS = ${NUM_CLASSIFICATIONS} +NUM_USERS = ${NUM_USERS} +NUM_TEAMS = ${NUM_TEAMS} +NUM_DOMAINS = ${NUM_DOMAINS} +NUM_DATA_PRODUCTS = ${NUM_DATA_PRODUCTS} +NUM_API_COLLECTIONS = ${NUM_API_COLLECTIONS} +NUM_API_ENDPOINTS = ${NUM_API_ENDPOINTS} +NUM_LINEAGE_EDGES = ${NUM_LINEAGE_EDGES} +NUM_TEST_CASE_RESULTS = ${NUM_TEST_CASE_RESULTS} +NUM_ENTITY_REPORT_DATA = ${NUM_ENTITY_REPORT_DATA} +NUM_WEB_ANALYTIC_VIEWS = ${NUM_WEB_ANALYTIC_VIEWS} +NUM_WEB_ANALYTIC_ACTIVITY = ${NUM_WEB_ANALYTIC_ACTIVITY} +NUM_RAW_COST_ANALYSIS = ${NUM_RAW_COST_ANALYSIS} +NUM_AGG_COST_ANALYSIS = ${NUM_AGG_COST_ANALYSIS} + +ONLY_ENTITIES_RAW = "${ONLY_ENTITIES}" +OUTPUT_PATH_RAW = "${OUTPUT_PATH}" +AUTH_TOKEN_RAW = "${AUTH_TOKEN}" +SCALE_APPLIED = "${SCALE_APPLIED}" or "default" +RAMP_MODE = "${RAMP_MODE}" == "true" +RAMP_BATCH = ${RAMP_BATCH} +ADMIN_PORT_RAW = "${ADMIN_PORT}" + +# Auto-calculate database/schema counts +NUM_DATABASES = max(1, NUM_TABLES // 50000) +SCHEMAS_PER_DB = min(20, max(1, NUM_TABLES // (NUM_DATABASES * 5000))) if NUM_TABLES > 0 else 1 + +# ── Entity filter (--only) ─────────────────────────────────────────────────── +ONLY_ENTITIES = set() +if ONLY_ENTITIES_RAW.strip(): + ONLY_ENTITIES = {e.strip() for e in ONLY_ENTITIES_RAW.split(",") if e.strip()} + +ENTITY_PREREQUISITES = { + "tables": {"_services_db", "_infra_db"}, + "storedProcedures": {"_services_db", "_infra_db"}, + "queries": {"_services_db"}, + "dashboards": {"_services_dashboard"}, + "charts": {"_services_dashboard"}, + "dashboardDataModels": {"_services_dashboard"}, + "topics": {"_services_messaging"}, + "pipelines": {"_services_pipeline"}, + "mlmodels": {"_services_mlmodel"}, + "containers": {"_services_storage"}, + "searchIndexes": {"_services_search"}, + "apiCollections": {"_services_api"}, + "apiEndpoints": {"_services_api", "apiCollections"}, + "dataProducts": {"domains"}, + "glossaryTerms": {"glossaries"}, + "tags": {"classifications"}, + "testCases": {"tables", "_services_db", "_infra_db"}, + "testCaseResults": {"testCases", "tables", "_services_db", "_infra_db"}, + "lineageEdges": {"tables", "_services_db", "_infra_db"}, +} + +def _resolve_prerequisites(entities): + resolved = set(entities) + changed = True + while changed: + changed = False + for e in list(resolved): + for prereq in ENTITY_PREREQUISITES.get(e, set()): + if prereq not in resolved: + resolved.add(prereq) + changed = True + return resolved + +if ONLY_ENTITIES: + _resolved = _resolve_prerequisites(ONLY_ENTITIES) + _infra_needed = {p for p in _resolved if p.startswith("_")} + _resolved -= _infra_needed +else: + _resolved = set() + _infra_needed = set() + +def should_run(entity_name): + if not ONLY_ENTITIES: + return True + return entity_name in ONLY_ENTITIES or entity_name in _resolved + +def _need_infra(tag): + if not ONLY_ENTITIES: + return True + return tag in _infra_needed + +print(f"Connecting to {SERVER_URL}...") +sys.stdout.flush() + +# ── HTTP helper with retry ─────────────────────────────────────────────────── +def make_request(url, data=None, method="GET", headers=None, retries=3): + if headers is None: + headers = {} + headers["Content-Type"] = "application/json" + encoded = json.dumps(data).encode("utf-8") if data else None + for attempt in range(retries): + req = urllib.request.Request(url, data=encoded, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=60) as resp: + body = resp.read().decode("utf-8") + try: + parsed = json.loads(body) if body.strip() else {} + except json.JSONDecodeError: + parsed = body + return resp.status, parsed + except urllib.error.HTTPError as e: + try: + body = e.read().decode("utf-8") + except Exception: + body = str(e) + if e.code >= 500 and attempt < retries - 1: + time.sleep(2 ** attempt) + continue + return e.code, body + except Exception as e: + if attempt < retries - 1: + time.sleep(2 ** attempt) + continue + return 0, str(e) + return 0, "max retries exceeded" + +# ── Auth token ─────────────────────────────────────────────────────────────── +if AUTH_TOKEN_RAW.strip(): + token = AUTH_TOKEN_RAW.strip() + print("Using token from --token argument.") +else: + token = ("eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg") + print("Using default admin JWT token for authentication.") + +headers = {"Content-Type": "application/json", "Authorization": f"Bearer {token}"} + +# ── Output path ────────────────────────────────────────────────────────────── +if OUTPUT_PATH_RAW.strip(): + output_path = OUTPUT_PATH_RAW.strip() +else: + ts_str = datetime.now().strftime("%Y%m%d-%H%M%S") + output_path = f"benchmark-report-{ts_str}.json" + +# ══════════════════════════════════════════════════════════════════════════════ +# SERVER INTROSPECTION +# ══════════════════════════════════════════════════════════════════════════════ + +class ServerIntrospector: + def __init__(self, server_url, req_headers, admin_port=None): + self.server_url = server_url + self.req_headers = req_headers + self.admin_port = admin_port + self.info = {} + + def _parse_url_host(self): + from urllib.parse import urlparse + parsed = urlparse(self.server_url) + return parsed.hostname or "localhost" + + def collect(self): + print("Collecting server introspection data...") + sys.stdout.flush() + + status, resp = make_request( + f"{self.server_url}/api/v1/system/version", + method="GET", headers=self.req_headers, retries=2, + ) + if status == 200 and isinstance(resp, dict): + self.info["version"] = resp.get("version", "unknown") + self.info["revision"] = resp.get("revision", "unknown") + print(f" Server version: {self.info['version']} (rev: {self.info['revision']})") + else: + self.info["version"] = "unreachable" + self.info["revision"] = "unknown" + print(f" Could not fetch version (status={status})") + + status, resp = make_request( + f"{self.server_url}/api/v1/system/status", + method="GET", headers=self.req_headers, retries=2, + ) + if status == 200 and isinstance(resp, dict): + self.info["status"] = resp + healthy_components = sum(1 for v in resp.values() if isinstance(v, dict) and v.get("passed")) + total_components = sum(1 for v in resp.values() if isinstance(v, dict)) + print(f" Component health: {healthy_components}/{total_components} healthy") + else: + self.info["status"] = {"error": f"status={status}"} + print(f" Could not fetch status (status={status})") + + status, resp = make_request( + f"{self.server_url}/api/v1/system/entities/count", + method="GET", headers=self.req_headers, retries=2, + ) + if status == 200 and isinstance(resp, dict): + self.info["entity_counts"] = resp + total_entities = sum(v for v in resp.values() if isinstance(v, (int, float))) + print(f" Pre-existing entities: {total_entities:,}") + else: + self.info["entity_counts"] = {} + print(f" Could not fetch entity counts (status={status})") + + status, resp = make_request( + f"{self.server_url}/api/v1/system/services/count", + method="GET", headers=self.req_headers, retries=2, + ) + if status == 200 and isinstance(resp, dict): + self.info["service_counts"] = resp + else: + self.info["service_counts"] = {} + + if self.admin_port: + self._scrape_prometheus("before") + + self.collect_diagnostics("before") + + sys.stdout.flush() + return self.info + + def _scrape_prometheus(self, label): + host = self._parse_url_host() + prom_url = f"http://{host}:{self.admin_port}/prometheus" + print(f" Scraping Prometheus metrics ({label}) from {prom_url}...") + sys.stdout.flush() + try: + req = urllib.request.Request(prom_url, method="GET") + with urllib.request.urlopen(req, timeout=10) as resp: + body = resp.read().decode("utf-8") + metrics = self._parse_prometheus(body) + key = f"prometheus_{label}" + self.info[key] = metrics + heap_used = metrics.get("jvm_memory_bytes_used_heap", "N/A") + heap_max = metrics.get("jvm_memory_bytes_max_heap", "N/A") + if isinstance(heap_used, (int, float)) and isinstance(heap_max, (int, float)): + print(f" JVM heap: {heap_used / 1048576:.0f}MB / {heap_max / 1048576:.0f}MB") + thread_count = metrics.get("jvm_threads_current", "N/A") + print(f" JVM threads: {thread_count}") + db_active = metrics.get("hikaricp_connections_active", "N/A") + db_idle = metrics.get("hikaricp_connections_idle", "N/A") + db_total = metrics.get("hikaricp_connections_total", "N/A") + print(f" DB pool: active={db_active}, idle={db_idle}, total={db_total}") + except Exception as e: + self.info[f"prometheus_{label}"] = {"error": str(e)} + print(f" Failed to scrape Prometheus: {e}") + sys.stdout.flush() + + def collect_diagnostics(self, label): + diag_url = f"{self.server_url}/api/v1/system/diagnostics" + print(f" Collecting diagnostics ({label}) from {diag_url}...") + sys.stdout.flush() + try: + status, resp = make_request(diag_url, method="GET", headers=self.req_headers, retries=2) + if status == 200 and isinstance(resp, dict): + self.info[f"diagnostics_{label}"] = resp + jvm = resp.get("jvm", {}) + jetty = resp.get("jetty", {}) + db = resp.get("database", {}) + bulk = resp.get("bulk_executor", {}) + heap_pct = jvm.get("heap_usage_pct", "N/A") + print(f" JVM heap: {heap_pct}% used, GC pauses: {jvm.get('gc_pause_total_ms', 0)}ms") + print(f" Jetty: {jetty.get('threads_busy', '?')}/{jetty.get('threads_max', '?')} " + f"threads busy ({jetty.get('utilization_pct', 0)}%), " + f"queue: {jetty.get('queue_size', 0)}") + print(f" DB pool: {db.get('pool_active', '?')}/{db.get('pool_max', '?')} " + f"active ({db.get('pool_usage_pct', 0)}%), " + f"pending: {db.get('pool_pending', 0)}") + print(f" Bulk executor: queue {bulk.get('queue_depth', 0)}/{bulk.get('queue_capacity', '?')} " + f"({bulk.get('queue_usage_pct', 0)}%)") + return resp + else: + print(f" Diagnostics endpoint returned status={status} (may not be available)") + self.info[f"diagnostics_{label}"] = None + return None + except Exception as e: + print(f" Failed to collect diagnostics: {e}") + self.info[f"diagnostics_{label}"] = None + return None + + def scrape_after(self): + if self.admin_port: + self._scrape_prometheus("after") + + @staticmethod + def _parse_prometheus(text): + metrics = {} + target_prefixes = [ + "jvm_memory_bytes_used", "jvm_memory_bytes_max", + "jvm_threads_current", "jvm_threads_daemon", + "hikaricp_connections_active", "hikaricp_connections_idle", + "hikaricp_connections_total", "hikaricp_connections_pending", + "io_dropwizard_jetty_MutableServletContextHandler_percent", + "io_dropwizard_jetty_MutableServletContextHandler_requests", + ] + for line in text.split("\n"): + if line.startswith("#") or not line.strip(): + continue + parts = line.split() + if len(parts) < 2: + continue + metric_name = parts[0] + try: + value = float(parts[1]) + except ValueError: + continue + for prefix in target_prefixes: + if metric_name.startswith(prefix): + if 'area="heap"' in metric_name: + if "bytes_used" in metric_name: + metrics["jvm_memory_bytes_used_heap"] = value + elif "bytes_max" in metric_name: + metrics["jvm_memory_bytes_max_heap"] = value + elif 'area="nonheap"' not in metric_name: + simple_name = metric_name.split("{")[0] + if simple_name not in metrics: + metrics[simple_name] = value + break + return metrics + + def report_section(self): + return dict(self.info) + + +introspector = ServerIntrospector( + SERVER_URL, dict(headers), + admin_port=ADMIN_PORT_RAW.strip() if ADMIN_PORT_RAW.strip() else None, +) +server_info = introspector.collect() + +# ══════════════════════════════════════════════════════════════════════════════ +# BENCHMARK INFRASTRUCTURE +# ══════════════════════════════════════════════════════════════════════════════ + +class BenchmarkCollector: + def __init__(self, entity_name, requested_count): + self.entity_name = entity_name + self.requested = requested_count + self.latencies = [] + self.errors = [] + self.created = 0 + self.failed = 0 + self.start_time = None + self.end_time = None + self.lock = threading.Lock() + self.window_counts = [] + + def record_success(self, latency_s): + with self.lock: + self.latencies.append(latency_s) + self.created += 1 + + def record_failure(self, latency_s, status, error): + with self.lock: + self.latencies.append(latency_s) + self.failed += 1 + if len(self.errors) < 50: + self.errors.append({"status": status, "error": str(error)[:200]}) + + def _error_breakdown(self): + breakdown = {} + for e in self.errors: + code = e.get("status", 0) + if code == 0: + key = "connection_error" + else: + key = str(code) + breakdown[key] = breakdown.get(key, 0) + 1 + return breakdown + + def _latency_analysis(self): + if len(self.latencies) < 10: + return {} + s = sorted(self.latencies) + n = len(s) + p50 = s[int(n * 0.50)] * 1000 + p90 = s[min(int(n * 0.90), n - 1)] * 1000 + p95 = s[min(int(n * 0.95), n - 1)] * 1000 + p99 = s[min(int(n * 0.99), n - 1)] * 1000 + + p90_p50_ratio = round(p90 / p50, 1) if p50 > 0 else 0 + p99_p95_ratio = round(p99 / p95, 1) if p95 > 0 else 0 + + bimodal = p90_p50_ratio > 5.0 + + degradation_pct = 0.0 + if n >= 20: + chunk = max(1, n // 5) + first_20 = s[:chunk] + last_20 = s[-chunk:] + avg_first = sum(first_20) / len(first_20) if first_20 else 0 + avg_last = sum(last_20) / len(last_20) if last_20 else 0 + if avg_first > 0: + degradation_pct = round((avg_last / avg_first - 1) * 100, 1) + + result = { + "bimodal": bimodal, + "p90_p50_ratio": p90_p50_ratio, + "p99_p95_ratio": p99_p95_ratio, + "degradation_pct": degradation_pct, + } + + findings = [] + if bimodal: + findings.append( + f"Bimodal latency distribution (p50={p50:.0f}ms, p90={p90:.0f}ms, " + f"ratio={p90_p50_ratio}x) -- likely DB connection pool or thread pool wait" + ) + if p99_p95_ratio > 3.0: + findings.append( + f"Extreme tail latency (p95={p95:.0f}ms, p99={p99:.0f}ms, " + f"ratio={p99_p95_ratio}x) -- possible GC pauses or lock contention" + ) + if degradation_pct > 100.0: + findings.append( + f"Sustained load degradation: last 20% of requests {degradation_pct:.0f}% " + f"slower than first 20% -- resource exhaustion" + ) + result["findings"] = findings + return result + + def percentile(self, p): + if not self.latencies: + return 0 + s = sorted(self.latencies) + k = int(len(s) * p / 100) + return s[min(k, len(s) - 1)] + + def _throughput_windows(self): + if not self.window_counts or len(self.window_counts) < 2: + return [] + buckets = [] + bucket_size = 5 + start_ts = self.window_counts[0][0] + i = 0 + while i < len(self.window_counts): + bucket_start = start_ts + len(buckets) * bucket_size + bucket_end = bucket_start + bucket_size + count_at_start = self.window_counts[i][1] if i == 0 else None + count_at_end = None + for j in range(i, len(self.window_counts)): + if self.window_counts[j][0] >= bucket_end: + break + count_at_end = self.window_counts[j][1] + if count_at_start is None: + count_at_start = self.window_counts[j][1] + i = j + 1 + if count_at_start is not None and count_at_end is not None: + delta = count_at_end - count_at_start + rps = delta / bucket_size if bucket_size > 0 else 0 + buckets.append({ + "elapsed_s": round(bucket_end - start_ts, 1), + "rps": round(rps, 1), + }) + else: + i += 1 + return buckets + + def summary(self): + n = len(self.latencies) + if n == 0: + return None + elapsed = (self.end_time or time.time()) - self.start_time if self.start_time else 0 + s = sorted(self.latencies) + result = { + "requested": self.requested, + "total_requests": n, + "created": self.created, + "failed": self.failed, + "error_rate_pct": round(self.failed / n * 100, 2) if n > 0 else 0, + "wall_clock_s": round(elapsed, 2), + "throughput_rps": round(n / elapsed, 2) if elapsed > 0 else 0, + "latency_ms": { + "min": round(s[0] * 1000, 1), + "p50": round(self.percentile(50) * 1000, 1), + "p75": round(self.percentile(75) * 1000, 1), + "p90": round(self.percentile(90) * 1000, 1), + "p95": round(self.percentile(95) * 1000, 1), + "p99": round(self.percentile(99) * 1000, 1), + "max": round(s[-1] * 1000, 1), + "avg": round(sum(s) / len(s) * 1000, 1), + }, + "throughput_over_time": self._throughput_windows(), + "errors_sample": self.errors[:10], + "latency_analysis": self._latency_analysis(), + "error_breakdown": self._error_breakdown(), + } + return result + + +class HealthMonitor: + def __init__(self, server_url, req_headers, interval=5, diagnostics_interval=10): + self.server_url = server_url + self.req_headers = req_headers + self.interval = interval + self.diagnostics_interval = diagnostics_interval + self.samples = [] + self.diagnostics_samples = [] + self.running = True + self.start_time = time.time() + self._last_diag_time = 0 + self._thread = threading.Thread(target=self._poll, daemon=True) + self._thread.start() + + def _poll(self): + while self.running: + t0 = time.time() + try: + status, _ = make_request( + f"{self.server_url}/api/v1/system/version", + method="GET", headers=self.req_headers, retries=1, + ) + except Exception: + status = 0 + latency = (time.time() - t0) * 1000 + self.samples.append({ + "timestamp": time.time(), + "elapsed_s": round(time.time() - self.start_time, 1), + "latency_ms": round(latency, 1), + "status": status, + "healthy": status == 200, + }) + now = time.time() + if now - self._last_diag_time >= self.diagnostics_interval: + self._last_diag_time = now + self._sample_diagnostics() + time.sleep(self.interval) + + def _sample_diagnostics(self): + diag_url = f"{self.server_url}/api/v1/system/diagnostics" + try: + status, resp = make_request(diag_url, method="GET", headers=self.req_headers, retries=1) + if status == 200 and isinstance(resp, dict): + resp["_sample_elapsed_s"] = round(time.time() - self.start_time, 1) + self.diagnostics_samples.append(resp) + except Exception: + pass + + def stop(self): + self.running = False + self._thread.join(timeout=self.interval + 2) + + def summary(self): + if not self.samples: + return { + "total_checks": 0, "healthy": 0, "unhealthy": 0, + "health_latency_ms": {}, "timeline": [], + } + healthy = [s for s in self.samples if s["healthy"]] + unhealthy = [s for s in self.samples if not s["healthy"]] + latencies = sorted([s["latency_ms"] for s in self.samples]) + n = len(latencies) + result = { + "total_checks": n, + "healthy": len(healthy), + "unhealthy": len(unhealthy), + "health_latency_ms": { + "min": round(latencies[0], 1), + "avg": round(sum(latencies) / n, 1), + "p95": round(latencies[min(int(n * 0.95), n - 1)], 1), + "max": round(latencies[-1], 1), + }, + "timeline": self.samples, + } + if self.diagnostics_samples: + result["diagnostics_during"] = self.diagnostics_samples + return result + + +# ── Start health monitor ──────────────────────────────────────────────────── +health_monitor = HealthMonitor(SERVER_URL, dict(headers), interval=5) + +overall_start = time.time() +stats = {} +benchmarks = {} +phase_timings = {} + +# ── Collected IDs for linking phases ───────────────────────────────────────── +collected_table_ids = [] +collected_dashboard_ids = [] +collected_pipeline_ids = [] +collected_test_case_fqns = [] +collect_lock = threading.Lock() + +MAX_COLLECT = max(NUM_LINEAGE_EDGES * 2, NUM_TEST_CASE_RESULTS, 10000) + +# ── Generic batch creator (with benchmarking) ─────────────────────────────── +def create_entity_batch(entity_name, count, payload_fn, workers=None, collect_fn=None, + log_interval=None): + if count <= 0: + return 0, None + if workers is None: + workers = NUM_WORKERS + if log_interval is None: + log_interval = max(1, count // 20) + + bench = BenchmarkCollector(entity_name, count) + bench.start_time = time.time() + + print(f"\nCreating {count} {entity_name}...") + sys.stdout.flush() + + counter_lock = threading.Lock() + sample_running = True + + def _sampler(): + while sample_running: + with counter_lock: + total = bench.created + bench.failed + bench.window_counts.append((time.time(), total)) + time.sleep(1) + + sampler_thread = threading.Thread(target=_sampler, daemon=True) + sampler_thread.start() + + def _work(idx): + t0 = time.time() + payload = payload_fn(idx) + if payload is None: + return + url = payload.pop("__url__", None) + method = payload.pop("__method__", "PUT") + if url is None: + return + status, resp = make_request(url, data=payload, method=method, headers=headers) + latency = time.time() - t0 + ok = status in [200, 201] + if ok: + bench.record_success(latency) + else: + bench.record_failure(latency, status, resp) + with counter_lock: + total = bench.created + bench.failed + if total % log_interval == 0 or total == count: + elapsed = time.time() - bench.start_time + rate = total / elapsed if elapsed > 0 else 0 + print(f" {entity_name}: {total}/{count} ({rate:.1f}/sec) - OK: {bench.created}, Fail: {bench.failed}") + sys.stdout.flush() + if ok and collect_fn and isinstance(resp, dict): + collect_fn(idx, resp) + + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = [executor.submit(_work, i) for i in range(count)] + for f in as_completed(futures): + try: + f.result() + except Exception: + pass + + sample_running = False + bench.end_time = time.time() + stats[entity_name] = bench.created + benchmarks[entity_name] = bench + elapsed = time.time() - bench.start_time + p50 = bench.percentile(50) * 1000 + p95 = bench.percentile(95) * 1000 + print(f"{entity_name} done: {bench.created} created, {bench.failed} failed " + f"({elapsed:.1f}s, p50={p50:.0f}ms, p95={p95:.0f}ms)") + sys.stdout.flush() + return bench.created, bench + + +# ── Ramp tester ────────────────────────────────────────────────────────────── +class RampTester: + def __init__(self, server_url, req_headers, max_workers, batch_size=100): + self.server_url = server_url + self.req_headers = req_headers + self.max_workers = max_workers + self.batch_size = batch_size + self.levels = [] + self.optimal_workers = 1 + self.saturation_point = None + + def _ramp_levels(self): + levels = [] + w = 1 + while w <= self.max_workers: + levels.append(w) + w *= 2 + if levels and levels[-1] != self.max_workers and self.max_workers > levels[-1]: + levels.append(self.max_workers) + return levels + + def run(self, entity_type="tables", schema_fqns_list=None, service_fqn=None): + levels = self._ramp_levels() + print(f"\nRAMP TEST ({entity_type}, {self.batch_size} entities per level)") + print(f"{'Workers':>7} {'RPS':>7} {'p50ms':>7} {'p95ms':>7} {'p99ms':>7} {'Errors':>6}") + sys.stdout.flush() + + best_rps = 0 + best_workers = 1 + prev_p95 = 0 + + for worker_count in levels: + level_result = self._run_level(worker_count, entity_type, schema_fqns_list, service_fqn) + self.levels.append(level_result) + + marker = "" + if level_result["rps"] > best_rps: + best_rps = level_result["rps"] + best_workers = worker_count + if level_result["rps"] >= best_rps: + self.optimal_workers = worker_count + + if (prev_p95 > 0 and level_result["p95_ms"] > prev_p95 * 3 + and self.saturation_point is None): + self.saturation_point = worker_count + marker = " <- saturation" + elif level_result["rps"] >= best_rps and level_result["errors"] == 0: + marker = " <- optimal" + + prev_p95 = level_result["p95_ms"] if level_result["p95_ms"] > 0 else prev_p95 + + print(f"{worker_count:>7} {level_result['rps']:>7.1f} {level_result['p50_ms']:>7.0f} " + f"{level_result['p95_ms']:>7.0f} {level_result['p99_ms']:>7.0f} " + f"{level_result['errors']:>6}{marker}") + sys.stdout.flush() + + self.optimal_workers = best_workers + print(f"Optimal concurrency: {self.optimal_workers} workers " + f"({best_rps:.1f} rps, p95={self._get_p95_for(self.optimal_workers):.0f}ms)") + if self.saturation_point: + print(f"Saturation point: {self.saturation_point} workers") + print("") + sys.stdout.flush() + + def _get_p95_for(self, workers): + for level in self.levels: + if level["workers"] == workers: + return level["p95_ms"] + return 0 + + def _run_level(self, worker_count, entity_type, schema_fqns_list, service_fqn): + latencies = [] + errors = 0 + error_lock = threading.Lock() + ramp_suffix = f"_ramp_{worker_count}w" + + def _payload_fn(idx): + if entity_type == "tables" and schema_fqns_list: + sfqn = schema_fqns_list[idx % len(schema_fqns_list)] if schema_fqns_list else "default" + return { + "__url__": f"{self.server_url}/api/v1/tables", + "name": f"ramp_table_{worker_count}w_{idx:07d}", + "databaseSchema": sfqn, + "columns": [ + {"name": "id", "dataType": "BIGINT"}, + {"name": "name", "dataType": "VARCHAR", "dataLength": 255}, + ], + } + elif entity_type == "topics" and service_fqn: + return { + "__url__": f"{self.server_url}/api/v1/topics", + "name": f"ramp_topic_{worker_count}w_{idx:06d}", + "service": service_fqn, + "partitions": 3, + "replicationFactor": 1, + } + else: + if schema_fqns_list: + sfqn = schema_fqns_list[idx % len(schema_fqns_list)] + return { + "__url__": f"{self.server_url}/api/v1/tables", + "name": f"ramp_table_{worker_count}w_{idx:07d}", + "databaseSchema": sfqn, + "columns": [ + {"name": "id", "dataType": "BIGINT"}, + {"name": "name", "dataType": "VARCHAR", "dataLength": 255}, + ], + } + return None + + def _work(idx): + nonlocal errors + t0 = time.time() + payload = _payload_fn(idx) + if payload is None: + return + url = payload.pop("__url__", None) + method = payload.pop("__method__", "PUT") + if url is None: + return + status, resp = make_request(url, data=payload, method=method, + headers=self.req_headers) + latency = time.time() - t0 + with error_lock: + latencies.append(latency) + if status not in [200, 201]: + errors += 1 + + start = time.time() + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [executor.submit(_work, i) for i in range(self.batch_size)] + for f in as_completed(futures): + try: + f.result() + except Exception: + with error_lock: + errors += 1 + elapsed = time.time() - start + + if not latencies: + return {"workers": worker_count, "rps": 0, "p50_ms": 0, "p95_ms": 0, + "p99_ms": 0, "errors": errors} + + s = sorted(latencies) + n = len(s) + return { + "workers": worker_count, + "rps": round(n / elapsed, 1) if elapsed > 0 else 0, + "p50_ms": round(s[int(n * 0.50)] * 1000, 0), + "p95_ms": round(s[min(int(n * 0.95), n - 1)] * 1000, 0), + "p99_ms": round(s[min(int(n * 0.99), n - 1)] * 1000, 0), + "errors": errors, + } + + def report_section(self): + if not self.levels: + return None + analysis_parts = [] + if self.optimal_workers: + opt = next((l for l in self.levels if l["workers"] == self.optimal_workers), None) + if opt: + analysis_parts.append( + f"Throughput peaks at {self.optimal_workers} workers ({opt['rps']} rps)." + ) + if self.saturation_point: + sat = next((l for l in self.levels if l["workers"] == self.saturation_point), None) + if sat: + analysis_parts.append( + f"At {self.saturation_point}+ workers, p95 exceeds {sat['p95_ms']:.0f}ms " + f"and errors appear -- thread pool saturation." + ) + return { + "batch_size": self.batch_size, + "levels": self.levels, + "optimal_workers": self.optimal_workers, + "saturation_point": self.saturation_point, + "analysis": " ".join(analysis_parts) if analysis_parts else "Ramp test completed.", + } + + +# ── Helper to create a service ─────────────────────────────────────────────── +def create_service(endpoint, data): + status, resp = make_request(f"{SERVER_URL}/api/v1/services/{endpoint}", data=data, + method="PUT", headers=headers) + if status in [200, 201] and isinstance(resp, dict): + fqn = resp["fullyQualifiedName"] + print(f" Service created: {fqn}") + return fqn + print(f" Failed to create service: {status} - {resp}") + return None + + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 1: Metadata (no dependencies) +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 1: Metadata (domains, classifications, glossaries, users, teams)") +print("=" * 60) + +# ── Domains ────────────────────────────────────────────────────────────────── +domain_fqns = [] + +if should_run("domains") and NUM_DOMAINS > 0: + def domain_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/domains", + "name": f"TestDomain_{idx:04d}", + "domainType": "Aggregate", + "description": f"Test domain {idx} for load testing", + } + + def collect_domain(idx, resp): + with collect_lock: + domain_fqns.append(resp.get("fullyQualifiedName", f"TestDomain_{idx:04d}")) + + create_entity_batch("domains", NUM_DOMAINS, domain_payload, collect_fn=collect_domain) + +# ── Classifications & Tags ─────────────────────────────────────────────────── +classification_fqns = [] +if should_run("classifications") and NUM_CLASSIFICATIONS > 0: + print(f"\nCreating {NUM_CLASSIFICATIONS} classifications...") + sys.stdout.flush() + for i in range(NUM_CLASSIFICATIONS): + data = { + "name": f"TestClassification_{i:04d}", + "description": f"Test classification {i} for load testing", + } + status, resp = make_request(f"{SERVER_URL}/api/v1/classifications", data=data, + method="PUT", headers=headers) + if status in [200, 201] and isinstance(resp, dict): + classification_fqns.append(resp["fullyQualifiedName"]) + stats["classifications"] = len(classification_fqns) + print(f"classifications done: {len(classification_fqns)} created") + +if should_run("tags") and NUM_TAGS > 0 and classification_fqns: + tags_per_class = max(1, NUM_TAGS // len(classification_fqns)) + tag_assignments = [] + for cfqn in classification_fqns: + for j in range(tags_per_class): + tag_assignments.append((cfqn, j)) + if len(tag_assignments) >= NUM_TAGS: + break + if len(tag_assignments) >= NUM_TAGS: + break + + def tag_payload(idx): + cfqn, j = tag_assignments[idx] + return { + "__url__": f"{SERVER_URL}/api/v1/tags", + "name": f"Tag_{j:05d}", + "classification": cfqn, + "description": f"Test tag {j}", + } + + create_entity_batch("tags", len(tag_assignments), tag_payload) + +# ── Glossaries & Terms ─────────────────────────────────────────────────────── +glossary_fqns = [] +if should_run("glossaries") and NUM_GLOSSARIES > 0: + print(f"\nCreating {NUM_GLOSSARIES} glossaries...") + sys.stdout.flush() + for i in range(NUM_GLOSSARIES): + data = { + "name": f"TestGlossary_{i:04d}", + "displayName": f"Test Glossary {i}", + "description": f"Test glossary {i} for load testing", + } + status, resp = make_request(f"{SERVER_URL}/api/v1/glossaries", data=data, + method="PUT", headers=headers) + if status in [200, 201] and isinstance(resp, dict): + glossary_fqns.append(resp["fullyQualifiedName"]) + stats["glossaries"] = len(glossary_fqns) + print(f"glossaries done: {len(glossary_fqns)} created") + +if should_run("glossaryTerms") and NUM_GLOSSARY_TERMS > 0 and glossary_fqns: + terms_per_glossary = max(1, NUM_GLOSSARY_TERMS // len(glossary_fqns)) + term_assignments = [] + for gfqn in glossary_fqns: + for j in range(terms_per_glossary): + term_assignments.append((gfqn, j)) + if len(term_assignments) >= NUM_GLOSSARY_TERMS: + break + if len(term_assignments) >= NUM_GLOSSARY_TERMS: + break + + def term_payload(idx): + gfqn, j = term_assignments[idx] + return { + "__url__": f"{SERVER_URL}/api/v1/glossaryTerms", + "name": f"Term_{j:05d}", + "glossary": gfqn, + "displayName": f"Term {j}", + "description": f"Test glossary term {j}", + } + + create_entity_batch("glossaryTerms", len(term_assignments), term_payload) + +# ── Users ──────────────────────────────────────────────────────────────────── +if should_run("users") and NUM_USERS > 0: + def user_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/users", + "name": f"testuser_{idx:05d}", + "email": f"testuser_{idx:05d}@example.com", + "displayName": f"Test User {idx}", + "description": f"Test user {idx} for load testing", + } + + create_entity_batch("users", NUM_USERS, user_payload) + +# ── Teams ──────────────────────────────────────────────────────────────────── +if should_run("teams") and NUM_TEAMS > 0: + def team_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/teams", + "name": f"testteam_{idx:04d}", + "displayName": f"Test Team {idx}", + "description": f"Test team {idx} for load testing", + "teamType": "Group", + } + + create_entity_batch("teams", NUM_TEAMS, team_payload) + +phase_timings["phase_1_metadata"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 2: Services +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 2: Services") +print("=" * 60) + +db_service_fqn = None +dashboard_service_fqn = None +pipeline_service_fqn = None +messaging_service_fqn = None +mlmodel_service_fqn = None +storage_service_fqn = None +search_service_fqn = None +api_service_fqn = None + +need_db_svc = (should_run("tables") or should_run("storedProcedures") or should_run("queries") + or should_run("testCases") or should_run("testCaseResults") + or should_run("lineageEdges") or _need_infra("_services_db")) +need_dashboard_svc = (should_run("dashboards") or should_run("charts") + or should_run("dashboardDataModels") or _need_infra("_services_dashboard")) +need_pipeline_svc = should_run("pipelines") or _need_infra("_services_pipeline") +need_messaging_svc = should_run("topics") or _need_infra("_services_messaging") +need_mlmodel_svc = should_run("mlmodels") or _need_infra("_services_mlmodel") +need_storage_svc = should_run("containers") or _need_infra("_services_storage") +need_search_svc = should_run("searchIndexes") or _need_infra("_services_search") +need_api_svc = (should_run("apiCollections") or should_run("apiEndpoints") + or _need_infra("_services_api")) + +if need_db_svc: + db_service_fqn = create_service("databaseServices", { + "name": "test-service-distributed", + "serviceType": "Mysql", + "connection": {"config": {"type": "Mysql", "username": "test", + "authType": {"password": "test"}, "hostPort": "localhost:3306"}}, + }) + +if need_dashboard_svc: + dashboard_service_fqn = create_service("dashboardServices", { + "name": "test-dashboard-service", + "serviceType": "Looker", + "connection": {"config": {"type": "Looker", "clientId": "test-client-id", + "clientSecret": "test-client-secret", + "hostPort": "https://looker.example.com"}}, + }) + +if need_pipeline_svc: + pipeline_service_fqn = create_service("pipelineServices", { + "name": "test-pipeline-service", + "serviceType": "Airflow", + "connection": {"config": {"type": "Airflow", "hostPort": "http://airflow.example.com:8080", + "connection": {"type": "Backend"}}}, + }) + +if need_messaging_svc: + messaging_service_fqn = create_service("messagingServices", { + "name": "test-messaging-service", + "serviceType": "Kafka", + "connection": {"config": {"type": "Kafka", "bootstrapServers": "localhost:9092"}}, + }) + +if need_mlmodel_svc: + mlmodel_service_fqn = create_service("mlmodelServices", { + "name": "test-mlmodel-service", + "serviceType": "Mlflow", + "connection": {"config": {"type": "Mlflow", "trackingUri": "http://mlflow.example.com:5000", + "registryUri": "http://mlflow.example.com:5000"}}, + }) + +if need_storage_svc: + storage_service_fqn = create_service("storageServices", { + "name": "test-storage-service", + "serviceType": "S3", + "connection": {"config": {"type": "S3", "awsConfig": { + "awsAccessKeyId": "test-key", "awsSecretAccessKey": "test-secret", + "awsRegion": "us-east-1"}}}, + }) + +if need_search_svc: + search_service_fqn = create_service("searchServices", { + "name": "test-search-service", + "serviceType": "ElasticSearch", + "connection": {"config": {"type": "ElasticSearch", + "hostPort": "http://elasticsearch.example.com:9200"}}, + }) + +if need_api_svc: + api_service_fqn = create_service("apiServices", { + "name": "test-api-service", + "serviceType": "Rest", + "connection": {"config": {"type": "Rest", + "openAPISchemaConnection": { + "openAPISchemaURL": "http://api.example.com/openapi.json"}}}, + }) + +phase_timings["phase_2_services"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 3: Infrastructure (databases, schemas, apiCollections) +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 3: Infrastructure (databases, schemas, API collections)") +print("=" * 60) + +# ── Databases & Schemas ────────────────────────────────────────────────────── +schema_fqns = [] +need_db_infra = (should_run("tables") or should_run("storedProcedures") or should_run("queries") + or should_run("testCases") or should_run("testCaseResults") + or should_run("lineageEdges") or _need_infra("_infra_db")) + +if db_service_fqn and need_db_infra and NUM_TABLES > 0: + print(f"Creating {NUM_DATABASES} databases with {SCHEMAS_PER_DB} schemas each...") + sys.stdout.flush() + for i in range(NUM_DATABASES): + db_data = {"name": f"test_db_{i:04d}", "service": db_service_fqn} + status, resp = make_request(f"{SERVER_URL}/api/v1/databases", data=db_data, + method="PUT", headers=headers) + if status in [200, 201] and isinstance(resp, dict): + db_fqn = resp["fullyQualifiedName"] + for s in range(SCHEMAS_PER_DB): + schema_name = f"schema_{s:03d}" if SCHEMAS_PER_DB > 1 else "public" + s_data = {"name": schema_name, "database": db_fqn} + s_status, s_resp = make_request(f"{SERVER_URL}/api/v1/databaseSchemas", + data=s_data, method="PUT", headers=headers) + if s_status in [200, 201] and isinstance(s_resp, dict): + schema_fqns.append(s_resp["fullyQualifiedName"]) + print(f" Created {NUM_DATABASES} databases, {len(schema_fqns)} schemas") + +# ── API Collections ────────────────────────────────────────────────────────── +api_collection_fqns = [] +if should_run("apiCollections") and api_service_fqn and NUM_API_COLLECTIONS > 0: + def api_coll_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/apiCollections", + "name": f"api_collection_{idx:04d}", + "service": api_service_fqn, + "displayName": f"API Collection {idx}", + "description": f"Test API collection {idx}", + } + + def collect_api_coll(idx, resp): + with collect_lock: + api_collection_fqns.append(resp.get("fullyQualifiedName", + f"test-api-service.api_collection_{idx:04d}")) + + create_entity_batch("apiCollections", NUM_API_COLLECTIONS, api_coll_payload, + collect_fn=collect_api_coll) + +phase_timings["phase_3_infrastructure"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# RAMP TEST (if enabled) +# ══════════════════════════════════════════════════════════════════════════════ +ramp_result = None +if RAMP_MODE: + ramp_tester = RampTester(SERVER_URL, dict(headers), max_workers=NUM_WORKERS, + batch_size=RAMP_BATCH) + ramp_entity_type = "tables" + if ONLY_ENTITIES: + for candidate in ["tables", "topics", "dashboards", "pipelines"]: + if candidate in ONLY_ENTITIES: + ramp_entity_type = candidate + break + + ramp_schema_fqns = schema_fqns if schema_fqns else None + ramp_service_fqn = None + if ramp_entity_type == "topics": + ramp_service_fqn = messaging_service_fqn + elif ramp_entity_type == "dashboards": + ramp_service_fqn = dashboard_service_fqn + elif ramp_entity_type == "pipelines": + ramp_service_fqn = pipeline_service_fqn + + ramp_tester.run(entity_type=ramp_entity_type, schema_fqns_list=ramp_schema_fqns, + service_fqn=ramp_service_fqn) + ramp_result = ramp_tester.report_section() + if ramp_result: + ramp_result["entity_type"] = ramp_entity_type + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 4: Core entities +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 4: Core entities") +print("=" * 60) + +# ── Tables ─────────────────────────────────────────────────────────────────── +if should_run("tables") and schema_fqns and NUM_TABLES > 0: + def table_payload(idx): + sfqn = schema_fqns[idx % len(schema_fqns)] + return { + "__url__": f"{SERVER_URL}/api/v1/tables", + "name": f"table_{idx:07d}", + "databaseSchema": sfqn, + "columns": [ + {"name": "id", "dataType": "BIGINT", "description": "Primary key"}, + {"name": "name", "dataType": "VARCHAR", "dataLength": 255}, + {"name": "created_at", "dataType": "TIMESTAMP"}, + {"name": "data", "dataType": "JSON"}, + ], + } + + def collect_table(idx, resp): + with collect_lock: + if len(collected_table_ids) < MAX_COLLECT: + collected_table_ids.append((resp["id"], resp.get("fullyQualifiedName", ""))) + + create_entity_batch("tables", NUM_TABLES, table_payload, collect_fn=collect_table) + +# ── Dashboards ─────────────────────────────────────────────────────────────── +if should_run("dashboards") and dashboard_service_fqn and NUM_DASHBOARDS > 0: + def dashboard_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/dashboards", + "name": f"dashboard_{idx:06d}", + "service": dashboard_service_fqn, + "displayName": f"Test Dashboard {idx}", + "description": f"Test dashboard {idx}", + } + + def collect_dashboard(idx, resp): + with collect_lock: + if len(collected_dashboard_ids) < MAX_COLLECT: + collected_dashboard_ids.append((resp["id"], resp.get("fullyQualifiedName", ""))) + + create_entity_batch("dashboards", NUM_DASHBOARDS, dashboard_payload, + collect_fn=collect_dashboard) + +# ── Charts ─────────────────────────────────────────────────────────────────── +if should_run("charts") and dashboard_service_fqn and NUM_CHARTS > 0: + chart_types = ["Line", "Bar", "Pie", "Area", "Scatter", "Table"] + + def chart_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/charts", + "name": f"chart_{idx:06d}", + "service": dashboard_service_fqn, + "displayName": f"Test Chart {idx}", + "chartType": chart_types[idx % len(chart_types)], + "description": f"Test chart {idx}", + } + + create_entity_batch("charts", NUM_CHARTS, chart_payload) + +# ── Topics ─────────────────────────────────────────────────────────────────── +if should_run("topics") and messaging_service_fqn and NUM_TOPICS > 0: + def topic_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/topics", + "name": f"topic_{idx:06d}", + "service": messaging_service_fqn, + "partitions": 3, + "replicationFactor": 1, + "description": f"Test topic {idx}", + } + + create_entity_batch("topics", NUM_TOPICS, topic_payload) + +# ── Pipelines ──────────────────────────────────────────────────────────────── +if should_run("pipelines") and pipeline_service_fqn and NUM_PIPELINES > 0: + def pipeline_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/pipelines", + "name": f"pipeline_{idx:06d}", + "service": pipeline_service_fqn, + "displayName": f"Test Pipeline {idx}", + "description": f"Test pipeline {idx}", + } + + def collect_pipeline(idx, resp): + with collect_lock: + if len(collected_pipeline_ids) < MAX_COLLECT: + collected_pipeline_ids.append((resp["id"], resp.get("fullyQualifiedName", ""))) + + create_entity_batch("pipelines", NUM_PIPELINES, pipeline_payload, + collect_fn=collect_pipeline) + +# ── Stored Procedures ──────────────────────────────────────────────────────── +if should_run("storedProcedures") and db_service_fqn and schema_fqns and NUM_STORED_PROCEDURES > 0: + def stored_proc_payload(idx): + sfqn = schema_fqns[idx % len(schema_fqns)] + return { + "__url__": f"{SERVER_URL}/api/v1/storedProcedures", + "name": f"stored_proc_{idx:06d}", + "databaseSchema": sfqn, + "storedProcedureCode": { + "code": f"CREATE PROCEDURE sp_{idx}() BEGIN SELECT 1; END", + "language": "SQL", + }, + "description": f"Test stored procedure {idx}", + } + + create_entity_batch("storedProcedures", NUM_STORED_PROCEDURES, stored_proc_payload) + +# ── ML Models ──────────────────────────────────────────────────────────────── +if should_run("mlmodels") and mlmodel_service_fqn and NUM_MLMODELS > 0: + algorithms = ["LinearRegression", "RandomForest", "XGBoost", "NeuralNetwork", + "SVM", "KMeans", "DecisionTree"] + + def mlmodel_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/mlmodels", + "name": f"mlmodel_{idx:06d}", + "service": mlmodel_service_fqn, + "algorithm": algorithms[idx % len(algorithms)], + "displayName": f"Test ML Model {idx}", + "description": f"Test ML model {idx}", + } + + create_entity_batch("mlmodels", NUM_MLMODELS, mlmodel_payload) + +# ── Containers ─────────────────────────────────────────────────────────────── +if should_run("containers") and storage_service_fqn and NUM_CONTAINERS > 0: + def container_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/containers", + "name": f"container_{idx:06d}", + "service": storage_service_fqn, + "displayName": f"Test Container {idx}", + "description": f"Test container {idx}", + } + + create_entity_batch("containers", NUM_CONTAINERS, container_payload) + +# ── Search Indexes ─────────────────────────────────────────────────────────── +if should_run("searchIndexes") and search_service_fqn and NUM_SEARCH_INDEXES > 0: + def search_idx_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/searchIndexes", + "name": f"search_index_{idx:06d}", + "service": search_service_fqn, + "displayName": f"Test Search Index {idx}", + "description": f"Test search index {idx}", + "fields": [ + {"name": "id", "dataType": "KEYWORD"}, + {"name": "content", "dataType": "TEXT"}, + {"name": "timestamp", "dataType": "DATE"}, + ], + } + + create_entity_batch("searchIndexes", NUM_SEARCH_INDEXES, search_idx_payload) + +# ── Queries ────────────────────────────────────────────────────────────────── +if should_run("queries") and NUM_QUERIES > 0 and db_service_fqn: + def query_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/queries", + "name": f"query_{idx:06d}", + "query": f"SELECT * FROM table_{idx % max(1, NUM_TABLES):07d} WHERE id = {idx}", + "service": db_service_fqn, + "description": f"Test query {idx}", + } + + create_entity_batch("queries", NUM_QUERIES, query_payload) + +# ── Dashboard Data Models ──────────────────────────────────────────────────── +if should_run("dashboardDataModels") and dashboard_service_fqn and NUM_DASHBOARD_DATA_MODELS > 0: + dm_types = ["MetabaseDataModel", "SupersetDataModel", "TableauDataModel"] + + def data_model_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/dashboard/datamodels", + "name": f"data_model_{idx:06d}", + "service": dashboard_service_fqn, + "dataModelType": dm_types[idx % len(dm_types)], + "displayName": f"Test Data Model {idx}", + "description": f"Test dashboard data model {idx}", + "columns": [ + {"name": "id", "dataType": "BIGINT"}, + {"name": "value", "dataType": "VARCHAR", "dataLength": 255}, + ], + } + + create_entity_batch("dashboardDataModels", NUM_DASHBOARD_DATA_MODELS, data_model_payload) + +# ── API Endpoints ──────────────────────────────────────────────────────────── +if should_run("apiEndpoints") and api_service_fqn and api_collection_fqns and NUM_API_ENDPOINTS > 0: + http_methods = ["GET", "POST", "PUT", "DELETE", "PATCH"] + + def api_endpoint_payload(idx): + coll_fqn = api_collection_fqns[idx % len(api_collection_fqns)] + return { + "__url__": f"{SERVER_URL}/api/v1/apiEndpoints", + "name": f"endpoint_{idx:06d}", + "apiCollection": coll_fqn, + "endpointURL": f"https://api.example.com/v1/resource_{idx}", + "requestMethod": http_methods[idx % len(http_methods)], + "displayName": f"Test Endpoint {idx}", + "description": f"Test API endpoint {idx}", + } + + create_entity_batch("apiEndpoints", NUM_API_ENDPOINTS, api_endpoint_payload) + +# ── Data Products ──────────────────────────────────────────────────────────── +if should_run("dataProducts") and NUM_DATA_PRODUCTS > 0 and domain_fqns: + def data_product_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/dataProducts", + "name": f"data_product_{idx:04d}", + "domains": [domain_fqns[idx % len(domain_fqns)]], + "displayName": f"Test Data Product {idx}", + "description": f"Test data product {idx}", + } + + create_entity_batch("dataProducts", NUM_DATA_PRODUCTS, data_product_payload) + +phase_timings["phase_4_core_entities"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 5: Data Quality (test suites, test cases) +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 5: Data Quality (test suites, test cases)") +print("=" * 60) + +test_suite_fqns = [] + +if should_run("testSuites") and NUM_TEST_SUITES > 0: + def test_suite_payload(idx): + return { + "__url__": f"{SERVER_URL}/api/v1/dataQuality/testSuites", + "name": f"testSuite_{idx:04d}", + "displayName": f"Test Suite {idx}", + "description": f"Test suite {idx} for load testing", + } + + def collect_test_suite(idx, resp): + with collect_lock: + test_suite_fqns.append(resp.get("fullyQualifiedName", + f"testSuite_{idx:04d}")) + + create_entity_batch("testSuites", NUM_TEST_SUITES, test_suite_payload, + collect_fn=collect_test_suite) + +if should_run("testCases") and NUM_TEST_CASES > 0 and collected_table_ids: + table_level_defs = [ + ("tableRowCountToEqual", [{"name": "value", "value": "100"}]), + ("tableColumnCountToEqual", [{"name": "columnCount", "value": "4"}]), + ] + column_level_defs = [ + ("columnValuesToBeNotNull", []), + ("columnValuesToBeUnique", []), + ] + column_names = ["id", "name", "created_at", "data"] + + def test_case_payload(idx): + table_id, table_fqn = collected_table_ids[idx % len(collected_table_ids)] + if idx % 3 == 0: + defn, params = table_level_defs[idx % len(table_level_defs)] + entity_link = f"<#E::table::{table_fqn}>" + else: + defn, params = column_level_defs[idx % len(column_level_defs)] + col = column_names[idx % len(column_names)] + entity_link = f"<#E::table::{table_fqn}::columns::{col}>" + return { + "__url__": f"{SERVER_URL}/api/v1/dataQuality/testCases", + "name": f"testCase_{idx:06d}", + "entityLink": entity_link, + "testDefinition": defn, + "parameterValues": params, + } + + def collect_test_case(idx, resp): + with collect_lock: + if len(collected_test_case_fqns) < MAX_COLLECT: + fqn = resp.get("fullyQualifiedName") + if fqn: + collected_test_case_fqns.append(fqn) + + create_entity_batch("testCases", NUM_TEST_CASES, test_case_payload, + collect_fn=collect_test_case) + +phase_timings["phase_5_data_quality"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 6: Lineage +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 6: Lineage") +print("=" * 60) + +if should_run("lineageEdges") and NUM_LINEAGE_EDGES > 0 and collected_table_ids: + n_t2t = int(NUM_LINEAGE_EDGES * 0.60) + n_t2d = int(NUM_LINEAGE_EDGES * 0.25) + n_p2t = NUM_LINEAGE_EDGES - n_t2t - n_t2d + + lineage_tasks = [] + for i in range(n_t2t): + from_idx = i % len(collected_table_ids) + to_idx = (i + 1) % len(collected_table_ids) + lineage_tasks.append(("table", collected_table_ids[from_idx][0], + "table", collected_table_ids[to_idx][0])) + if collected_dashboard_ids: + for i in range(n_t2d): + from_idx = i % len(collected_table_ids) + to_idx = i % len(collected_dashboard_ids) + lineage_tasks.append(("table", collected_table_ids[from_idx][0], + "dashboard", collected_dashboard_ids[to_idx][0])) + else: + n_t2t += n_t2d + n_t2d = 0 + + if collected_pipeline_ids: + for i in range(n_p2t): + from_idx = i % len(collected_pipeline_ids) + to_idx = (i + len(collected_table_ids) // 2) % len(collected_table_ids) + lineage_tasks.append(("pipeline", collected_pipeline_ids[from_idx][0], + "table", collected_table_ids[to_idx][0])) + else: + n_t2t += n_p2t + n_p2t = 0 + + actual_edges = len(lineage_tasks) + + def lineage_payload(idx): + if idx >= len(lineage_tasks): + return None + from_type, from_id, to_type, to_id = lineage_tasks[idx] + return { + "__url__": f"{SERVER_URL}/api/v1/lineage", + "__method__": "PUT", + "edge": { + "fromEntity": {"id": from_id, "type": from_type}, + "toEntity": {"id": to_id, "type": to_type}, + }, + } + + create_entity_batch("lineageEdges", actual_edges, lineage_payload, + workers=min(10, NUM_WORKERS)) +elif should_run("lineageEdges") and NUM_LINEAGE_EDGES > 0: + print("Skipping lineage: no table IDs collected") + stats["lineageEdges"] = 0 + +phase_timings["phase_6_lineage"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# PHASE 7: Time-Series Data +# ══════════════════════════════════════════════════════════════════════════════ +phase_start = time.time() +print("\n" + "=" * 60) +print("PHASE 7: Time-Series Data") +print("=" * 60) + +ts_workers = min(10, NUM_WORKERS) +base_ts = int(time.time() * 1000) + +# ── Test Case Results ──────────────────────────────────────────────────────── +if should_run("testCaseResults") and NUM_TEST_CASE_RESULTS > 0 and collected_test_case_fqns: + statuses = ["Success", "Failed", "Aborted"] + + def tc_result_payload(idx): + fqn = collected_test_case_fqns[idx % len(collected_test_case_fqns)] + encoded_fqn = urllib.request.quote(fqn, safe="") + ts = base_ts - (idx * 3600000) + return { + "__url__": f"{SERVER_URL}/api/v1/dataQuality/testCases/testCaseResults/{encoded_fqn}", + "__method__": "POST", + "timestamp": ts, + "testCaseStatus": statuses[idx % len(statuses)], + "result": f"Test result {idx}", + "testResultValue": [{"name": "value", "value": str(round(random.uniform(0, 100), 2))}], + } + + create_entity_batch("testCaseResults", NUM_TEST_CASE_RESULTS, tc_result_payload, + workers=ts_workers) +elif should_run("testCaseResults") and NUM_TEST_CASE_RESULTS > 0: + print("Skipping testCaseResults: no test case FQNs collected") + stats["testCaseResults"] = 0 + +# ── Entity Report Data ─────────────────────────────────────────────────────── +if should_run("entityReportData") and NUM_ENTITY_REPORT_DATA > 0: + entity_types_for_report = ["table", "topic", "dashboard", "pipeline", "mlmodel"] + + def entity_report_payload(idx): + ts = base_ts - (idx * 86400000) + e_type = entity_types_for_report[idx % len(entity_types_for_report)] + entity_count = random.randint(1, 1000) + has_owner = random.randint(0, entity_count) + return { + "__url__": f"{SERVER_URL}/api/v1/analytics/dataInsights/data", + "__method__": "POST", + "timestamp": ts, + "reportDataType": "entityReportData", + "data": { + "entityType": e_type, + "entityTier": f"Tier.Tier{(idx % 5) + 1}", + "serviceName": "test-service-distributed", + "completedDescriptions": random.randint(0, 100), + "missingDescriptions": random.randint(0, 50), + "hasOwner": has_owner, + "missingOwner": entity_count - has_owner, + "entityCount": entity_count, + }, + } + + create_entity_batch("entityReportData", NUM_ENTITY_REPORT_DATA, entity_report_payload, + workers=ts_workers) + +# ── Web Analytic Entity Views ──────────────────────────────────────────────── +if should_run("webAnalyticViews") and NUM_WEB_ANALYTIC_VIEWS > 0: + def web_view_payload(idx): + ts = base_ts - (idx * 60000) + return { + "__url__": f"{SERVER_URL}/api/v1/analytics/dataInsights/data", + "__method__": "POST", + "timestamp": ts, + "reportDataType": "webAnalyticEntityViewReportData", + "data": { + "entityType": "table", + "entityFqn": f"test-service-distributed.test_db_0000.public.table_{idx % max(1, NUM_TABLES):07d}", + "entityHref": f"{SERVER_URL}/table/test-service-distributed.test_db_0000.public.table_{idx % max(1, NUM_TABLES):07d}", + "owner": f"user_{idx % 50}", + "views": random.randint(1, 500), + }, + } + + create_entity_batch("webAnalyticViews", NUM_WEB_ANALYTIC_VIEWS, web_view_payload, + workers=ts_workers) + +# ── Web Analytic User Activity ─────────────────────────────────────────────── +if should_run("webAnalyticActivity") and NUM_WEB_ANALYTIC_ACTIVITY > 0: + def web_activity_payload(idx): + ts = base_ts - (idx * 60000) + return { + "__url__": f"{SERVER_URL}/api/v1/analytics/dataInsights/data", + "__method__": "POST", + "timestamp": ts, + "reportDataType": "webAnalyticUserActivityReportData", + "data": { + "userName": f"testuser_{idx % max(1, NUM_USERS):05d}", + "userId": str(uuid.uuid4()), + "team": "test-team", + "totalSessions": random.randint(1, 20), + "totalSessionDuration": random.randint(10, 3600), + "totalPageView": random.randint(1, 100), + "lastSession": ts, + }, + } + + create_entity_batch("webAnalyticActivity", NUM_WEB_ANALYTIC_ACTIVITY, + web_activity_payload, workers=ts_workers) + +# ── Raw Cost Analysis ──────────────────────────────────────────────────────── +if should_run("rawCostAnalysis") and NUM_RAW_COST_ANALYSIS > 0: + def raw_cost_payload(idx): + ts = base_ts - (idx * 86400000) + if collected_table_ids: + table_id, table_fqn = collected_table_ids[idx % len(collected_table_ids)] + else: + table_id = str(uuid.uuid4()) + table_fqn = f"test-service-distributed.test_db_0000.public.table_{idx % max(1, NUM_TABLES):07d}" + return { + "__url__": f"{SERVER_URL}/api/v1/analytics/dataInsights/data", + "__method__": "POST", + "timestamp": ts, + "reportDataType": "rawCostAnalysisReportData", + "data": { + "entity": { + "id": table_id, + "type": "table", + "fullyQualifiedName": table_fqn, + }, + "sizeInByte": round(random.uniform(100.0, 100000.0), 2), + }, + } + + create_entity_batch("rawCostAnalysis", NUM_RAW_COST_ANALYSIS, raw_cost_payload, + workers=ts_workers) + +# ── Aggregated Cost Analysis ───────────────────────────────────────────────── +if should_run("aggCostAnalysis") and NUM_AGG_COST_ANALYSIS > 0: + def agg_cost_payload(idx): + ts = base_ts - (idx * 86400000) + return { + "__url__": f"{SERVER_URL}/api/v1/analytics/dataInsights/data", + "__method__": "POST", + "timestamp": ts, + "reportDataType": "aggregatedCostAnalysisReportData", + "data": { + "entityType": "table", + "serviceName": "test-service-distributed", + "serviceType": "BigQuery", + "totalSize": round(random.uniform(1000.0, 1000000.0), 2), + "totalCount": random.randint(100, 100000), + "unusedDataAssets": { + "count": {"threeDays": random.randint(1, 50), "sevenDays": random.randint(1, 40), + "fourteenDays": random.randint(1, 30), "thirtyDays": random.randint(1, 20), + "sixtyDays": random.randint(1, 10)}, + "size": {"threeDays": random.randint(100, 10000), "sevenDays": random.randint(100, 9000), + "fourteenDays": random.randint(100, 8000), "thirtyDays": random.randint(100, 7000), + "sixtyDays": random.randint(100, 6000)}, + "totalSize": random.randint(1000, 50000), + "totalCount": random.randint(1, 50), + }, + "frequentlyUsedDataAssets": { + "count": {"threeDays": random.randint(1, 10), "sevenDays": random.randint(1, 20), + "fourteenDays": random.randint(1, 30), "thirtyDays": random.randint(1, 40), + "sixtyDays": random.randint(1, 50)}, + "size": {"threeDays": random.randint(1000, 50000), "sevenDays": random.randint(1000, 60000), + "fourteenDays": random.randint(1000, 70000), "thirtyDays": random.randint(1000, 80000), + "sixtyDays": random.randint(1000, 90000)}, + "totalSize": random.randint(1000, 100000), + "totalCount": random.randint(1, 50), + }, + }, + } + + create_entity_batch("aggCostAnalysis", NUM_AGG_COST_ANALYSIS, agg_cost_payload, + workers=ts_workers) + +phase_timings["phase_7_time_series"] = {"wall_clock_s": round(time.time() - phase_start, 2)} + +# ══════════════════════════════════════════════════════════════════════════════ +# STOP HEALTH MONITOR & BUILD REPORT +# ══════════════════════════════════════════════════════════════════════════════ +health_monitor.stop() +introspector.scrape_after() +introspector.collect_diagnostics("after") +overall_elapsed = time.time() - overall_start +total_created = sum(stats.values()) + + +# ── Cluster sizing analysis ────────────────────────────────────────────────── +def generate_cluster_sizing(report, ramp_data=None): + findings = [] + recommendations = {} + config_summary = [] + + entity_data = report.get("entities", {}) + health = report.get("server_health", {}) + measured_workers = report["metadata"]["workers"] + + max_p95 = 0 + avg_latency_ms = 0 + total_latency_entries = 0 + has_503 = False + has_504_or_timeout = False + has_conn_refused = False + has_bimodal = False + + for entity, data in entity_data.items(): + if not data: + continue + p95 = data.get("latency_ms", {}).get("p95", 0) + avg = data.get("latency_ms", {}).get("avg", 0) + if p95 > max_p95: + max_p95 = p95 + total_latency_entries += 1 + avg_latency_ms += avg + if p95 > 5000: + findings.append(f"{entity} PUT p95={p95:.0f}ms -- severe bottleneck") + elif p95 > 1000: + findings.append(f"{entity} PUT p95={p95:.0f}ms -- moderate latency") + + latency_analysis = data.get("latency_analysis", {}) + if latency_analysis.get("bimodal"): + has_bimodal = True + for finding in latency_analysis.get("findings", []): + findings.append(f"{entity}: {finding}") + + error_breakdown = data.get("error_breakdown", {}) + for code, count in error_breakdown.items(): + if code == "503": + has_503 = True + findings.append( + f"{entity}: {count}x HTTP 503 -- admission control rejection " + f"(BulkExecutor queue full)" + ) + elif code in ("504", "timeout"): + has_504_or_timeout = True + findings.append( + f"{entity}: {count}x HTTP {code} -- request timeout, " + f"thread pool likely exhausted" + ) + elif code == "connection_error": + has_conn_refused = True + findings.append( + f"{entity}: {count}x connection refused/reset -- " + f"server accept queue full" + ) + elif code == "500": + findings.append(f"{entity}: {count}x HTTP 500 -- internal server error, check logs") + elif code == "429": + findings.append(f"{entity}: {count}x HTTP 429 -- rate limited") + + if total_latency_entries > 0: + avg_latency_ms = avg_latency_ms / total_latency_entries + + total_checks = health.get("total_checks", 0) + unhealthy_count = health.get("unhealthy", 0) + if unhealthy_count > 0 and total_checks > 0: + pct = unhealthy_count / total_checks * 100 + findings.append(f"Health check failed {unhealthy_count} times ({pct:.0f}%) -- server overwhelmed") + + for entity, data in entity_data.items(): + if not data: + continue + windows = data.get("throughput_over_time", []) + if len(windows) >= 4: + q_len = max(1, len(windows) // 4) + first_q = [w["rps"] for w in windows[:q_len]] + last_q = [w["rps"] for w in windows[-q_len:]] + if first_q and last_q: + avg_first = sum(first_q) / len(first_q) + avg_last = sum(last_q) / len(last_q) + if avg_first > 0 and avg_last / avg_first < 0.5: + degradation = (1 - avg_last / avg_first) * 100 + findings.append( + f"{entity}: throughput degraded {degradation:.0f}% " + f"({avg_first:.0f} rps -> {avg_last:.0f} rps) -- resource exhaustion" + ) + + overall_rps = report.get("overall", {}).get("overall_throughput_rps", 0) + effective_capacity = int(150 / (avg_latency_ms / 1000)) if avg_latency_ms > 0 else 0 + + recommendations["server_threads"] = { + "analysis": ( + f"{measured_workers} concurrent writers with 150 max server threads. " + f"Each PUT holds a thread for ~{avg_latency_ms:.0f}ms avg. " + f"Effective capacity: ~{effective_capacity} rps. " + f"Measured: {overall_rps:.0f} rps" + + (" -- threads are blocked on DB/search." if overall_rps < effective_capacity * 0.5 else ".") + ), + "current_env": "SERVER_MAX_THREADS=150", + "recommended_env": "SERVER_MAX_THREADS=300" if max_p95 > 2000 or has_504_or_timeout else "SERVER_MAX_THREADS=150", + "yaml_path": "server.applicationConnectors[0].maxThreads", + } + if max_p95 > 2000 or has_504_or_timeout: + config_summary.append("SERVER_MAX_THREADS=300") + + recommendations["virtual_threads"] = { + "analysis": "Java 21 virtual threads can eliminate thread pool as bottleneck for I/O-bound workloads.", + "current_env": "SERVER_ENABLE_VIRTUAL_THREAD=false", + "recommended_env": "SERVER_ENABLE_VIRTUAL_THREAD=true", + "yaml_path": "server.enableVirtualThreads", + } + config_summary.append("SERVER_ENABLE_VIRTUAL_THREAD=true") + + db_pool_rec = 100 + if measured_workers > 15 or has_bimodal: + db_pool_rec = 150 + if max_p95 > 5000: + db_pool_rec = 200 + recommendations["db_pool"] = { + "analysis": ( + f"{measured_workers} concurrent writers with 100 max DB connections. " + f"Each PUT does multiple DB round-trips. " + f"At {measured_workers} concurrent, pool utilization ~{min(99, measured_workers * 100 // 100)}%. " + f"Connection wait time may add 50-200ms." + ), + "current_env": "DB_CONNECTION_POOL_MAX_SIZE=100", + "recommended_env": f"DB_CONNECTION_POOL_MAX_SIZE={db_pool_rec}", + "yaml_path": "database.hikari.maximumPoolSize", + } + if db_pool_rec > 100: + config_summary.append(f"DB_CONNECTION_POOL_MAX_SIZE={db_pool_rec}") + + if has_bimodal: + recommendations["db_connection_timeout"] = { + "analysis": ( + "Connection timeout 30s. If pool is exhausted, requests wait up to 30s " + "for a connection -- this explains the bimodal latency pattern." + ), + "current_env": "DB_CONNECTION_TIMEOUT=30000", + "recommended_env": "DB_CONNECTION_TIMEOUT=10000", + "note": "Fail faster to prevent cascading timeouts", + } + config_summary.append("DB_CONNECTION_TIMEOUT=10000") + + if max_p95 > 2000 or measured_workers > 10: + recommendations["search_connections"] = { + "analysis": ( + "Search max connections: 30 total, 10 per route. " + "Async indexing means search rarely blocks PUTs directly, " + "but under heavy write load the search queue can back up." + ), + "current_env": "ELASTICSEARCH_MAX_CONN_TOTAL=30", + "recommended_env": "ELASTICSEARCH_MAX_CONN_TOTAL=50", + } + config_summary.append("ELASTICSEARCH_MAX_CONN_TOTAL=50") + + if has_503: + recommendations["bulk_operation"] = { + "analysis": "HTTP 503 errors indicate BulkExecutor queue is full.", + "current_env": "BULK_OPERATION_QUEUE_SIZE=1000, BULK_OPERATION_MAX_THREADS=10", + "recommended_env": "BULK_OPERATION_QUEUE_SIZE=2000, BULK_OPERATION_MAX_THREADS=20", + } + config_summary.append("BULK_OPERATION_QUEUE_SIZE=2000") + config_summary.append("BULK_OPERATION_MAX_THREADS=20") + + if has_conn_refused: + recommendations["accept_queue"] = { + "analysis": "Connection refused/reset errors indicate server accept queue is full.", + "current_env": "SERVER_ACCEPT_QUEUE_SIZE=256", + "recommended_env": "SERVER_ACCEPT_QUEUE_SIZE=512", + } + config_summary.append("SERVER_ACCEPT_QUEUE_SIZE=512") + + ramp_optimal = None + if ramp_data: + ramp_optimal = ramp_data.get("optimal_workers", measured_workers) + recommendations["api_concurrency"] = { + "current": measured_workers, + "recommended": ramp_optimal, + "reason": f"Ramp test shows optimal throughput at {ramp_optimal} workers", + } + else: + if max_p95 > 5000: + rec_workers = max(2, measured_workers // 4) + elif max_p95 > 2000: + rec_workers = max(2, measured_workers // 2) + else: + rec_workers = measured_workers + recommendations["api_concurrency"] = { + "current": measured_workers, + "recommended": rec_workers, + "reason": f"Max p95 latency {max_p95:.0f}ms across entity types", + } + + if overall_rps > 0: + rec_rps = overall_rps + if ramp_optimal and ramp_optimal != measured_workers: + scale_factor = ramp_optimal / measured_workers if measured_workers > 0 else 1 + rec_rps = overall_rps * max(0.5, min(2.0, scale_factor)) + estimates = {} + for target in [50000, 250000, 1000000, 5000000]: + secs = target / rec_rps + label = f"{target // 1000}k_entities" + if secs < 3600: + estimates[label] = f"~{secs / 60:.0f} min" + else: + estimates[label] = f"~{secs / 3600:.1f} hours" + recommendations["estimated_times"] = estimates + + if max_p95 > 5000 or (unhealthy_count / max(1, total_checks)) > 0.1 or has_503: + assessment = "undersized" + elif max_p95 > 2000 or has_bimodal: + assessment = "marginal" + else: + assessment = "adequate" + + server_side = _build_server_side_analysis(report, findings, recommendations, config_summary) + + return { + "assessment": assessment, + "findings": findings, + "recommendations": recommendations, + "config_summary": config_summary, + "server_side_analysis": server_side, + } + + +def _build_server_side_analysis(report, findings, recommendations, config_summary): + si = report.get("server_info", {}) + diag_before = si.get("diagnostics_before") + diag_after = si.get("diagnostics_after") + if not diag_after: + return {"available": False} + + analysis = {"available": True, "bottlenecks": [], "latency_breakdown": {}} + + jvm = diag_after.get("jvm", {}) + jetty = diag_after.get("jetty", {}) + db = diag_after.get("database", {}) + bulk = diag_after.get("bulk_executor", {}) + req_latency = diag_after.get("request_latency", {}) + + analysis["snapshot_after"] = { + "jvm_heap_pct": jvm.get("heap_usage_pct", 0), + "gc_pause_total_ms": jvm.get("gc_pause_total_ms", 0), + "jetty_utilization_pct": jetty.get("utilization_pct", 0), + "jetty_queue_size": jetty.get("queue_size", 0), + "db_pool_usage_pct": db.get("pool_usage_pct", 0), + "db_pool_pending": db.get("pool_pending", 0), + "bulk_queue_usage_pct": bulk.get("queue_usage_pct", 0), + } + + if diag_before: + jvm_before = diag_before.get("jvm", {}) + gc_before = jvm_before.get("gc_pause_total_ms", 0) + gc_after = jvm.get("gc_pause_total_ms", 0) + analysis["gc_pause_delta_ms"] = gc_after - gc_before + + for ep_key, ep_data in req_latency.items(): + db_pct = ep_data.get("db_pct", 0) + search_pct = ep_data.get("search_pct", 0) + internal_pct = ep_data.get("internal_pct", 0) + db_pool_pct = db.get("pool_usage_pct", 0) + + if db_pct > 60 and db_pool_pct > 80: + bottleneck = (f"DB bottleneck on {ep_key}: {db_pct}% of request time in DB, " + f"pool at {db_pool_pct}% utilization") + analysis["bottlenecks"].append(bottleneck) + findings.append(bottleneck) + if "db_pool_increase" not in recommendations: + recommendations["db_pool_increase"] = { + "analysis": f"DB pool at {db_pool_pct}% with {db_pct}% of latency in DB", + "recommended_env": "DB_CONNECTION_POOL_MAX_SIZE=150", + } + config_summary.append("DB_CONNECTION_POOL_MAX_SIZE=150") + + if search_pct > 30: + bottleneck = f"Search pressure on {ep_key}: {search_pct}% of request time in search" + analysis["bottlenecks"].append(bottleneck) + findings.append(bottleneck) + + analysis["latency_breakdown"][ep_key] = { + "avg_total_ms": ep_data.get("avg_total_ms", 0), + "db_pct": db_pct, + "search_pct": search_pct, + "internal_pct": internal_pct, + } + + jetty_util = jetty.get("utilization_pct", 0) + jetty_queue = jetty.get("queue_size", 0) + if jetty_util > 90 and jetty_queue > 0: + queue_time = jetty.get("queue_time_avg_ms", 0) + bottleneck = (f"Thread pool saturated: {jetty_util}% utilization, " + f"{jetty_queue} requests queued, avg queue wait {queue_time}ms") + analysis["bottlenecks"].append(bottleneck) + findings.append(bottleneck) + + bulk_queue_pct = bulk.get("queue_usage_pct", 0) + if bulk_queue_pct > 70: + bottleneck = f"Bulk executor queue at {bulk_queue_pct}%, near rejection threshold" + analysis["bottlenecks"].append(bottleneck) + findings.append(bottleneck) + if "bulk_operation" not in recommendations: + recommendations["bulk_operation"] = { + "analysis": f"Bulk executor queue at {bulk_queue_pct}%", + "recommended_env": "BULK_OPERATION_QUEUE_SIZE=2000, BULK_OPERATION_MAX_THREADS=20", + } + config_summary.append("BULK_OPERATION_QUEUE_SIZE=2000") + config_summary.append("BULK_OPERATION_MAX_THREADS=20") + + heap_pct = jvm.get("heap_usage_pct", 0) + if heap_pct > 85: + bottleneck = f"JVM heap at {heap_pct}%, GC pressure likely causing tail latency" + analysis["bottlenecks"].append(bottleneck) + findings.append(bottleneck) + + return analysis + + +# ── Build report ───────────────────────────────────────────────────────────── +entity_summaries = {} +for name, bench in benchmarks.items(): + s = bench.summary() + if s: + entity_summaries[name] = s + +report = { + "metadata": { + "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + "server_url": SERVER_URL, + "workers": NUM_WORKERS, + "scale": SCALE_APPLIED, + "only_entities": list(ONLY_ENTITIES) if ONLY_ENTITIES else None, + "script_version": "3.0", + }, + "server_info": introspector.report_section(), + "server_health": health_monitor.summary(), + "entities": entity_summaries, + "phases": phase_timings, + "overall": { + "total_entities_created": total_created, + "total_wall_clock_s": round(overall_elapsed, 2), + "overall_throughput_rps": round(total_created / overall_elapsed, 2) if overall_elapsed > 0 else 0, + "total_errors": sum(b.failed for b in benchmarks.values()), + "overall_error_rate_pct": round( + sum(b.failed for b in benchmarks.values()) / + max(1, sum(len(b.latencies) for b in benchmarks.values())) * 100, 2 + ), + }, +} + +if ramp_result: + report["ramp_test"] = ramp_result + +si = introspector.report_section() +if si.get("diagnostics_before"): + report["diagnostics_before"] = si["diagnostics_before"] +if si.get("diagnostics_after"): + report["diagnostics_after"] = si["diagnostics_after"] + +report["cluster_sizing"] = generate_cluster_sizing(report, ramp_data=ramp_result) + +# ── Write JSON report ──────────────────────────────────────────────────────── +report_for_file = json.loads(json.dumps(report)) +if "server_health" in report_for_file and "timeline" in report_for_file["server_health"]: + del report_for_file["server_health"]["timeline"] + +try: + with open(output_path, "w") as f: + json.dump(report_for_file, f, indent=2) + print(f"\nJSON report written to: {output_path}") +except Exception as e: + print(f"\nFailed to write JSON report: {e}") + + +# ── Pretty-print summary table ────────────────────────────────────────────── +def print_summary_table(report): + print("") + print("\u2550" * 70) + print("BENCHMARK RESULTS") + print("\u2550" * 70) + + si = report.get("server_info", {}) + if si.get("version"): + print(f"Server: {si.get('version', '?')} (rev: {si.get('revision', '?')[:8]})") + print("") + + header = f"{'Entity':<22} {'Count':>7} {'Rate/s':>8} {'p50ms':>7} {'p95ms':>7} {'p99ms':>7} {'Errors':>7}" + print(header) + print("\u2500" * 70) + + for name, data in report["entities"].items(): + count = data["created"] + rps = data["throughput_rps"] + p50 = data["latency_ms"]["p50"] + p95 = data["latency_ms"]["p95"] + p99 = data["latency_ms"]["p99"] + err_pct = data["error_rate_pct"] + err_str = f"{err_pct:.1f}%" if err_pct > 0 else "0.0%" + print(f"{name:<22} {count:>7} {rps:>8.1f} {p50:>7.0f} {p95:>7.0f} {p99:>7.0f} {err_str:>7}") + + print("\u2500" * 70) + overall = report["overall"] + total = overall["total_entities_created"] + overall_rps = overall["overall_throughput_rps"] + overall_err = overall["overall_error_rate_pct"] + err_str = f"{overall_err:.1f}%" if overall_err > 0 else "0.0%" + print(f"{'Overall':<22} {total:>7} {overall_rps:>8.1f} {'':>7} {'':>7} {'':>7} {err_str:>7}") + print("") + + has_error_breakdowns = False + for name, data in report["entities"].items(): + eb = data.get("error_breakdown", {}) + if eb: + if not has_error_breakdowns: + print("ERROR BREAKDOWN:") + has_error_breakdowns = True + codes = ", ".join(f"{code}={cnt}" for code, cnt in sorted(eb.items())) + print(f" {name}: {codes}") + if has_error_breakdowns: + print("") + + has_latency_findings = False + for name, data in report["entities"].items(): + la = data.get("latency_analysis", {}) + for finding in la.get("findings", []): + if not has_latency_findings: + print("LATENCY ANALYSIS:") + has_latency_findings = True + print(f" {name}: {finding}") + if has_latency_findings: + print("") + + health = report["server_health"] + total_checks = health["total_checks"] + healthy = health["healthy"] + if total_checks > 0: + health_pct = healthy / total_checks * 100 + health_p95 = health.get("health_latency_ms", {}).get("p95", 0) + print(f"Server Health: {healthy}/{total_checks} checks healthy ({health_pct:.1f}%)") + print(f"Health p95 latency: {health_p95:.0f}ms") + else: + print("Server Health: no checks recorded") + print("") + + sizing = report.get("cluster_sizing", {}) + assessment = sizing.get("assessment", "unknown") + indicator = {"undersized": "!! UNDERSIZED", "marginal": "? MARGINAL", "adequate": "OK ADEQUATE"}.get( + assessment, f" {assessment.upper()}" + ) + print(f"CLUSTER SIZING: {indicator}") + for finding in sizing.get("findings", []): + print(f" - {finding}") + recs = sizing.get("recommendations", {}) + concurrency = recs.get("api_concurrency", {}) + if concurrency: + print(f" - Recommend: {concurrency.get('recommended', '?')} workers " + f"(currently {concurrency.get('current', '?')})") + estimates = recs.get("estimated_times", {}) + if "1000k_entities" in estimates: + print(f" - At current rate: 1M entities = {estimates['1000k_entities']}") + print("") + + ssa = sizing.get("server_side_analysis", {}) + if ssa.get("available"): + print("SERVER-SIDE BREAKDOWN (from /api/v1/system/diagnostics):") + snap = ssa.get("snapshot_after", {}) + gc_delta = ssa.get("gc_pause_delta_ms", 0) + si_info = report.get("server_info", {}) + diag_after = si_info.get("diagnostics_after", {}) + jvm_d = diag_after.get("jvm", {}) + jetty_d = diag_after.get("jetty", {}) + db_d = diag_after.get("database", {}) + bulk_d = diag_after.get("bulk_executor", {}) + + heap_used_gb = jvm_d.get("heap_used_bytes", 0) / (1024**3) + heap_max_gb = jvm_d.get("heap_max_bytes", 0) / (1024**3) + print(f" JVM: heap {heap_used_gb:.1f}GB/{heap_max_gb:.1f}GB " + f"({snap.get('jvm_heap_pct', 0)}%), " + f"GC pauses +{gc_delta}ms during load") + print(f" Jetty: {jetty_d.get('threads_busy', '?')}/{jetty_d.get('threads_max', '?')} " + f"threads busy ({snap.get('jetty_utilization_pct', 0)}%), " + f"queue depth: {snap.get('jetty_queue_size', 0)}") + print(f" DB Pool: {db_d.get('pool_active', '?')}/{db_d.get('pool_max', '?')} " + f"active ({snap.get('db_pool_usage_pct', 0)}%), " + f"{snap.get('db_pool_pending', 0)} pending connections") + print(f" Bulk Executor: queue {bulk_d.get('queue_depth', 0)}/" + f"{bulk_d.get('queue_capacity', '?')} " + f"({snap.get('bulk_queue_usage_pct', 0)}%)") + + breakdown = ssa.get("latency_breakdown", {}) + put_entries = {k: v for k, v in breakdown.items() if k.startswith("PUT")} + if put_entries: + print("") + print(" Latency Breakdown (PUT endpoints):") + print(f" {'Endpoint':<30} {'Total':>8} {'DB%':>6} {'Search%':>9} {'Internal%':>11}") + for ep, bd in sorted(put_entries.items()): + ep_short = ep.replace("PUT ", "")[:28] + print(f" {ep_short:<30} {bd['avg_total_ms']:>7.0f}ms " + f"{bd['db_pct']:>5.1f}% {bd['search_pct']:>8.1f}% " + f"{bd['internal_pct']:>10.1f}%") + + bottlenecks = ssa.get("bottlenecks", []) + if bottlenecks: + print("") + primary = bottlenecks[0] + print(f" BOTTLENECK: {primary}") + print("") + + config_summary = sizing.get("config_summary", []) + if config_summary: + print("RECOMMENDED CONFIGURATION:") + for env_line in config_summary: + print(f" export {env_line}") + print("") + + print(f"Full report: {output_path}") + print(f"Total time: {overall['total_wall_clock_s']:.1f}s") + +print_summary_table(report) + +print("") +print("Collected IDs for linking:") +print(f" Table IDs: {len(collected_table_ids)}") +print(f" Dashboard IDs: {len(collected_dashboard_ids)}") +print(f" Pipeline IDs: {len(collected_pipeline_ids)}") +print(f" Test Case FQNs: {len(collected_test_case_fqns)}") +PYEOF + +echo "" +echo "Test data loaded. You can now trigger reindexing with: ./scripts/trigger-reindex.sh" diff --git a/docker/development/distributed-test/scripts/start.sh b/bin/distributed-test/scripts/start.sh similarity index 97% rename from docker/development/distributed-test/scripts/start.sh rename to bin/distributed-test/scripts/start.sh index 423fa7f6c02..80fea78ae25 100755 --- a/docker/development/distributed-test/scripts/start.sh +++ b/bin/distributed-test/scripts/start.sh @@ -115,7 +115,7 @@ echo "All services are up and running!" echo "======================================" echo "" echo "Next steps:" -echo " 1. Load test data: ./scripts/load-test-data.sh --tables 10000" +echo " 1. Load test data: ./scripts/perf-test.sh --tables 10000" echo " 2. Trigger reindexing: ./scripts/trigger-reindex.sh" echo " 3. Watch logs: ./scripts/logs.sh -f" echo "" diff --git a/docker/development/distributed-test/scripts/stop.sh b/bin/distributed-test/scripts/stop.sh similarity index 100% rename from docker/development/distributed-test/scripts/stop.sh rename to bin/distributed-test/scripts/stop.sh diff --git a/docker/development/distributed-test/scripts/trigger-reindex.sh b/bin/distributed-test/scripts/trigger-reindex.sh similarity index 100% rename from docker/development/distributed-test/scripts/trigger-reindex.sh rename to bin/distributed-test/scripts/trigger-reindex.sh diff --git a/docker/development/distributed-test/README.md b/docker/development/distributed-test/README.md index 1213635e6a7..8237c6d2005 100644 --- a/docker/development/distributed-test/README.md +++ b/docker/development/distributed-test/README.md @@ -49,10 +49,10 @@ cd docker/development/distributed-test ```bash # Load 10,000 tables (default) -./scripts/load-test-data.sh +./scripts/perf-test.sh # Or specify the number -./scripts/load-test-data.sh --tables 50000 --databases 50 +./scripts/perf-test.sh --tables 50000 --databases 50 ``` ### 3. Trigger Reindexing @@ -163,7 +163,7 @@ OPENMETADATA_HEAP_OPTS=-Xmx1G -Xms1G ### Verify Partition Distribution 1. Start all 3 servers -2. Load test data: `./scripts/load-test-data.sh --tables 10000` +2. Load test data: `./scripts/perf-test.sh --tables 10000` 3. Trigger reindex: `./scripts/trigger-reindex.sh --recreate` 4. Watch logs: `./scripts/logs.sh -f --grep "partition"` @@ -242,7 +242,7 @@ distributed-test/ │ ├── stop.sh # Stop environment │ ├── logs.sh # View aggregated logs │ ├── trigger-reindex.sh # Trigger reindexing -│ └── load-test-data.sh # Load test data +│ └── perf-test.sh # Load test data ├── local/ │ ├── docker-compose-deps.yml # Dependencies only (for IDE debugging) │ ├── server1.yaml # Server 1 config (port 8585) diff --git a/docker/development/distributed-test/scripts/load-test-data.sh b/docker/development/distributed-test/scripts/load-test-data.sh deleted file mode 100755 index 29cbe81af5b..00000000000 --- a/docker/development/distributed-test/scripts/load-test-data.sh +++ /dev/null @@ -1,960 +0,0 @@ -#!/bin/bash -# Load test data for distributed indexing testing - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -# Default values -SERVER_URL="http://localhost:8585" -NUM_TABLES=20000 -NUM_DASHBOARDS=5000 -NUM_CHARTS=10000 -NUM_PIPELINES=3000 -NUM_TOPICS=3000 -NUM_MLMODELS=2000 -NUM_CONTAINERS=2000 -NUM_SEARCH_INDEXES=1000 -NUM_GLOSSARIES=50 -NUM_TERMS_PER_GLOSSARY=100 -NUM_CLASSIFICATIONS=20 -NUM_TAGS_PER_CLASSIFICATION=50 -NUM_DATABASES=10 - -# Parse arguments -while [[ $# -gt 0 ]]; do - case $1 in - --tables) - NUM_TABLES="$2" - shift 2 - ;; - --dashboards) - NUM_DASHBOARDS="$2" - shift 2 - ;; - --charts) - NUM_CHARTS="$2" - shift 2 - ;; - --pipelines) - NUM_PIPELINES="$2" - shift 2 - ;; - --topics) - NUM_TOPICS="$2" - shift 2 - ;; - --mlmodels) - NUM_MLMODELS="$2" - shift 2 - ;; - --containers) - NUM_CONTAINERS="$2" - shift 2 - ;; - --search-indexes) - NUM_SEARCH_INDEXES="$2" - shift 2 - ;; - --glossaries) - NUM_GLOSSARIES="$2" - shift 2 - ;; - --terms-per-glossary) - NUM_TERMS_PER_GLOSSARY="$2" - shift 2 - ;; - --classifications) - NUM_CLASSIFICATIONS="$2" - shift 2 - ;; - --tags-per-classification) - NUM_TAGS_PER_CLASSIFICATION="$2" - shift 2 - ;; - --databases) - NUM_DATABASES="$2" - shift 2 - ;; - --server) - SERVER_URL="$2" - shift 2 - ;; - --quick) - # Quick mode for rapid testing - smaller dataset - NUM_TABLES=3000 - NUM_DASHBOARDS=1000 - NUM_CHARTS=2000 - NUM_PIPELINES=500 - NUM_TOPICS=500 - NUM_MLMODELS=300 - NUM_CONTAINERS=300 - NUM_SEARCH_INDEXES=200 - NUM_GLOSSARIES=10 - NUM_TERMS_PER_GLOSSARY=50 - NUM_CLASSIFICATIONS=5 - NUM_TAGS_PER_CLASSIFICATION=20 - shift - ;; - -h|--help) - echo "Usage: $0 [OPTIONS]" - echo "" - echo "Options:" - echo " --tables NUM Number of tables to create (default: 20000)" - echo " --dashboards NUM Number of dashboards to create (default: 5000)" - echo " --charts NUM Number of charts to create (default: 10000)" - echo " --pipelines NUM Number of pipelines to create (default: 3000)" - echo " --topics NUM Number of topics to create (default: 3000)" - echo " --mlmodels NUM Number of ML models to create (default: 2000)" - echo " --containers NUM Number of containers to create (default: 2000)" - echo " --search-indexes NUM Number of search indexes to create (default: 1000)" - echo " --glossaries NUM Number of glossaries to create (default: 50)" - echo " --terms-per-glossary NUM Number of terms per glossary (default: 100)" - echo " --classifications NUM Number of classifications to create (default: 20)" - echo " --tags-per-classification Number of tags per classification (default: 50)" - echo " --databases NUM Number of databases (default: 10)" - echo " --server URL Target server URL (default: http://localhost:8585)" - echo " --quick Quick mode with smaller dataset (~10k entities)" - echo " -h, --help Show this help message" - exit 0 - ;; - *) - echo "Unknown option: $1" - exit 1 - ;; - esac -done - -NUM_GLOSSARY_TERMS=$((NUM_GLOSSARIES * NUM_TERMS_PER_GLOSSARY)) -NUM_TAGS=$((NUM_CLASSIFICATIONS * NUM_TAGS_PER_CLASSIFICATION)) -TOTAL=$((NUM_TABLES + NUM_DASHBOARDS + NUM_CHARTS + NUM_PIPELINES + NUM_TOPICS + NUM_MLMODELS + NUM_CONTAINERS + NUM_SEARCH_INDEXES + NUM_GLOSSARIES + NUM_GLOSSARY_TERMS + NUM_CLASSIFICATIONS + NUM_TAGS)) - -echo "======================================" -echo "Loading Test Data for Distributed Indexing" -echo "======================================" -echo "Server: $SERVER_URL" -echo "" -echo "Entity counts:" -echo " - Tables: $NUM_TABLES" -echo " - Dashboards: $NUM_DASHBOARDS" -echo " - Charts: $NUM_CHARTS" -echo " - Pipelines: $NUM_PIPELINES" -echo " - Topics: $NUM_TOPICS" -echo " - ML Models: $NUM_MLMODELS" -echo " - Containers: $NUM_CONTAINERS" -echo " - Search Indexes: $NUM_SEARCH_INDEXES" -echo " - Glossaries: $NUM_GLOSSARIES" -echo " - Glossary Terms: $NUM_GLOSSARY_TERMS" -echo " - Classifications: $NUM_CLASSIFICATIONS" -echo " - Tags: $NUM_TAGS" -echo " --------------------------" -echo " - Total: $TOTAL" -echo "" - -# Use Python with urllib (built-in, no extra packages needed) -python3 << EOF -import urllib.request -import urllib.error -import json -import sys -import time -import random -from concurrent.futures import ThreadPoolExecutor, as_completed - -SERVER_URL = "${SERVER_URL}" -NUM_TABLES = ${NUM_TABLES} -NUM_DASHBOARDS = ${NUM_DASHBOARDS} -NUM_CHARTS = ${NUM_CHARTS} -NUM_PIPELINES = ${NUM_PIPELINES} -NUM_TOPICS = ${NUM_TOPICS} -NUM_MLMODELS = ${NUM_MLMODELS} -NUM_CONTAINERS = ${NUM_CONTAINERS} -NUM_SEARCH_INDEXES = ${NUM_SEARCH_INDEXES} -NUM_GLOSSARIES = ${NUM_GLOSSARIES} -NUM_TERMS_PER_GLOSSARY = ${NUM_TERMS_PER_GLOSSARY} -NUM_CLASSIFICATIONS = ${NUM_CLASSIFICATIONS} -NUM_TAGS_PER_CLASSIFICATION = ${NUM_TAGS_PER_CLASSIFICATION} -NUM_DATABASES = ${NUM_DATABASES} - -print(f"Connecting to {SERVER_URL}...") -sys.stdout.flush() - -def make_request(url, data=None, method="GET", headers=None): - if headers is None: - headers = {} - headers["Content-Type"] = "application/json" - if data: - data = json.dumps(data).encode('utf-8') - req = urllib.request.Request(url, data=data, headers=headers, method=method) - try: - with urllib.request.urlopen(req, timeout=60) as response: - return response.status, json.loads(response.read().decode('utf-8')) - except urllib.error.HTTPError as e: - try: - body = e.read().decode('utf-8') - except: - body = str(e) - return e.code, body - except Exception as e: - return 0, str(e) - -# Use admin JWT token directly -token = "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg" -print("Using admin JWT token for authentication.") - -headers = {"Content-Type": "application/json"} -if token: - headers["Authorization"] = f"Bearer {token}" - -overall_start = time.time() -stats = { - "tables": 0, "dashboards": 0, "charts": 0, "pipelines": 0, "topics": 0, - "mlmodels": 0, "containers": 0, "searchIndexes": 0, - "glossaries": 0, "glossaryTerms": 0, "classifications": 0, "tags": 0 -} - -# Store created entity FQNs for linking -dashboard_fqns = [] - -# ========== CLASSIFICATIONS & TAGS ========== -if NUM_CLASSIFICATIONS > 0: - print("") - print("=" * 50) - print("Creating Classifications and Tags") - print("=" * 50) - - created_classifications = 0 - created_tags = 0 - start_time = time.time() - - for i in range(NUM_CLASSIFICATIONS): - classification_data = { - "name": f"TestClassification_{i:04d}", - "description": f"Test classification {i} for distributed indexing testing" - } - status, resp = make_request( - f"{SERVER_URL}/api/v1/classifications", - data=classification_data, - method="PUT", - headers=headers - ) - if status in [200, 201] and isinstance(resp, dict): - created_classifications += 1 - classification_fqn = resp["fullyQualifiedName"] - - # Create tags for this classification - def create_tag(args): - class_fqn, tag_idx = args - tag_data = { - "name": f"Tag_{tag_idx:04d}", - "classification": class_fqn, - "description": f"Test tag {tag_idx} in classification" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/tags", data=tag_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=10) as executor: - futures = [executor.submit(create_tag, (classification_fqn, j)) for j in range(NUM_TAGS_PER_CLASSIFICATION)] - for future in as_completed(futures): - if future.result(): - created_tags += 1 - - if (i + 1) % 5 == 0 or i == NUM_CLASSIFICATIONS - 1: - elapsed = time.time() - start_time - print(f" Classifications: {created_classifications}/{NUM_CLASSIFICATIONS}, Tags: {created_tags}/{NUM_CLASSIFICATIONS * NUM_TAGS_PER_CLASSIFICATION}") - sys.stdout.flush() - - stats["classifications"] = created_classifications - stats["tags"] = created_tags - print(f"Classifications completed: {created_classifications} created") - print(f"Tags completed: {created_tags} created") - -# ========== GLOSSARIES & TERMS ========== -if NUM_GLOSSARIES > 0: - print("") - print("=" * 50) - print("Creating Glossaries and Terms") - print("=" * 50) - - created_glossaries = 0 - created_terms = 0 - start_time = time.time() - - for i in range(NUM_GLOSSARIES): - glossary_data = { - "name": f"TestGlossary_{i:04d}", - "displayName": f"Test Glossary {i}", - "description": f"Test glossary {i} for distributed indexing testing" - } - status, resp = make_request( - f"{SERVER_URL}/api/v1/glossaries", - data=glossary_data, - method="PUT", - headers=headers - ) - if status in [200, 201] and isinstance(resp, dict): - created_glossaries += 1 - glossary_fqn = resp["fullyQualifiedName"] - - # Create terms for this glossary - def create_term(args): - gloss_fqn, term_idx = args - term_data = { - "name": f"Term_{term_idx:04d}", - "glossary": gloss_fqn, - "displayName": f"Term {term_idx}", - "description": f"Test glossary term {term_idx}" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/glossaryTerms", data=term_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=10) as executor: - futures = [executor.submit(create_term, (glossary_fqn, j)) for j in range(NUM_TERMS_PER_GLOSSARY)] - for future in as_completed(futures): - if future.result(): - created_terms += 1 - - if (i + 1) % 10 == 0 or i == NUM_GLOSSARIES - 1: - elapsed = time.time() - start_time - print(f" Glossaries: {created_glossaries}/{NUM_GLOSSARIES}, Terms: {created_terms}/{NUM_GLOSSARIES * NUM_TERMS_PER_GLOSSARY}") - sys.stdout.flush() - - stats["glossaries"] = created_glossaries - stats["glossaryTerms"] = created_terms - print(f"Glossaries completed: {created_glossaries} created") - print(f"Glossary Terms completed: {created_terms} created") - -# ========== TABLES ========== -if NUM_TABLES > 0: - print("") - print("=" * 50) - print("Creating Tables") - print("=" * 50) - - # Create database service - print("Creating database service...") - sys.stdout.flush() - service_data = { - "name": "test-service-distributed", - "serviceType": "Mysql", - "connection": { - "config": { - "type": "Mysql", - "username": "test", - "authType": {"password": "test"}, - "hostPort": "localhost:3306" - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/databaseServices", - data=service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - db_service_fqn = resp["fullyQualifiedName"] - print(f"Database service created: {db_service_fqn}") - else: - print(f"Failed to create database service: {status} - {resp}") - sys.exit(1) - - # Create databases - print(f"Creating {NUM_DATABASES} databases...") - sys.stdout.flush() - database_fqns = [] - for i in range(NUM_DATABASES): - db_data = {"name": f"test_db_{i:04d}", "service": db_service_fqn} - status, resp = make_request(f"{SERVER_URL}/api/v1/databases", data=db_data, method="PUT", headers=headers) - if status in [200, 201] and isinstance(resp, dict): - database_fqns.append(resp["fullyQualifiedName"]) - - print(f"Created {len(database_fqns)} databases") - - # Create schemas - print("Creating schemas...") - schema_fqns = [] - for db_fqn in database_fqns: - schema_data = {"name": "public", "database": db_fqn} - status, resp = make_request(f"{SERVER_URL}/api/v1/databaseSchemas", data=schema_data, method="PUT", headers=headers) - if status in [200, 201] and isinstance(resp, dict): - schema_fqns.append(resp["fullyQualifiedName"]) - - print(f"Created {len(schema_fqns)} schemas") - - if not schema_fqns: - print("ERROR: No schemas created. Cannot continue with tables.") - else: - # Create tables - print(f"Creating {NUM_TABLES} tables...") - sys.stdout.flush() - tables_per_schema = NUM_TABLES // len(schema_fqns) - created = 0 - failed = 0 - start_time = time.time() - - def create_table(args): - schema_fqn, table_idx = args - table_data = { - "name": f"table_{table_idx:06d}", - "databaseSchema": schema_fqn, - "columns": [ - {"name": "id", "dataType": "BIGINT", "description": "Primary key"}, - {"name": "name", "dataType": "VARCHAR", "dataLength": 255, "description": "Name field"}, - {"name": "created_at", "dataType": "TIMESTAMP", "description": "Creation timestamp"}, - {"name": "data", "dataType": "JSON", "description": "JSON data field"} - ] - } - status, _ = make_request(f"{SERVER_URL}/api/v1/tables", data=table_data, method="PUT", headers=headers) - return status in [200, 201] - - # Prepare tasks - tasks = [] - table_idx = 0 - for schema_fqn in schema_fqns: - for _ in range(tables_per_schema): - tasks.append((schema_fqn, table_idx)) - table_idx += 1 - if table_idx >= NUM_TABLES: - break - if table_idx >= NUM_TABLES: - break - - # Execute with thread pool - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_table, task): task for task in tasks} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 1000 == 0 or total == len(tasks): - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Tables: {total}/{NUM_TABLES} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["tables"] = created - print(f"Tables completed: {created} created, {failed} failed") - -# ========== DASHBOARDS ========== -if NUM_DASHBOARDS > 0: - print("") - print("=" * 50) - print("Creating Dashboards") - print("=" * 50) - - # Create dashboard service - print("Creating dashboard service...") - sys.stdout.flush() - dashboard_service_data = { - "name": "test-dashboard-service", - "serviceType": "Looker", - "connection": { - "config": { - "type": "Looker", - "clientId": "test-client-id", - "clientSecret": "test-client-secret", - "hostPort": "https://looker.example.com" - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/dashboardServices", - data=dashboard_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - dashboard_service_fqn = resp["fullyQualifiedName"] - print(f"Dashboard service created: {dashboard_service_fqn}") - else: - print(f"Failed to create dashboard service: {status} - {resp}") - dashboard_service_fqn = None - - if dashboard_service_fqn: - # Create dashboards - print(f"Creating {NUM_DASHBOARDS} dashboards...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_dashboard(idx): - dashboard_data = { - "name": f"dashboard_{idx:06d}", - "service": dashboard_service_fqn, - "displayName": f"Test Dashboard {idx}", - "description": f"Auto-generated test dashboard {idx} for distributed indexing testing" - } - status, resp = make_request(f"{SERVER_URL}/api/v1/dashboards", data=dashboard_data, method="PUT", headers=headers) - if status in [200, 201] and isinstance(resp, dict): - return resp.get("fullyQualifiedName") - return None - - # Execute with thread pool - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_dashboard, i): i for i in range(NUM_DASHBOARDS)} - for future in as_completed(futures): - result = future.result() - if result: - created += 1 - dashboard_fqns.append(result) - else: - failed += 1 - total = created + failed - if total % 1000 == 0 or total == NUM_DASHBOARDS: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Dashboards: {total}/{NUM_DASHBOARDS} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["dashboards"] = created - print(f"Dashboards completed: {created} created, {failed} failed") - -# ========== CHARTS ========== -if NUM_CHARTS > 0 and dashboard_fqns: - print("") - print("=" * 50) - print("Creating Charts") - print("=" * 50) - - print(f"Creating {NUM_CHARTS} charts linked to dashboards...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_chart(idx): - # Link chart to a random dashboard - dashboard_fqn = random.choice(dashboard_fqns) if dashboard_fqns else None - chart_data = { - "name": f"chart_{idx:06d}", - "service": dashboard_service_fqn, - "displayName": f"Test Chart {idx}", - "chartType": random.choice(["Line", "Bar", "Pie", "Area", "Scatter", "Table"]), - "description": f"Auto-generated test chart {idx}" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/charts", data=chart_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_chart, i): i for i in range(NUM_CHARTS)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 1000 == 0 or total == NUM_CHARTS: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Charts: {total}/{NUM_CHARTS} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["charts"] = created - print(f"Charts completed: {created} created, {failed} failed") - -# ========== PIPELINES ========== -if NUM_PIPELINES > 0: - print("") - print("=" * 50) - print("Creating Pipelines") - print("=" * 50) - - # Create pipeline service - print("Creating pipeline service...") - sys.stdout.flush() - pipeline_service_data = { - "name": "test-pipeline-service", - "serviceType": "Airflow", - "connection": { - "config": { - "type": "Airflow", - "hostPort": "http://airflow.example.com:8080", - "connection": { - "type": "BackendConnection" - } - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/pipelineServices", - data=pipeline_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - pipeline_service_fqn = resp["fullyQualifiedName"] - print(f"Pipeline service created: {pipeline_service_fqn}") - else: - print(f"Failed to create pipeline service: {status} - {resp}") - pipeline_service_fqn = None - - if pipeline_service_fqn: - # Create pipelines - print(f"Creating {NUM_PIPELINES} pipelines...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_pipeline(idx): - pipeline_data = { - "name": f"pipeline_{idx:06d}", - "service": pipeline_service_fqn, - "displayName": f"Test Pipeline {idx}", - "description": f"Auto-generated test pipeline {idx} for distributed indexing testing" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/pipelines", data=pipeline_data, method="PUT", headers=headers) - return status in [200, 201] - - # Execute with thread pool - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_pipeline, i): i for i in range(NUM_PIPELINES)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 1000 == 0 or total == NUM_PIPELINES: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Pipelines: {total}/{NUM_PIPELINES} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["pipelines"] = created - print(f"Pipelines completed: {created} created, {failed} failed") - -# ========== TOPICS ========== -if NUM_TOPICS > 0: - print("") - print("=" * 50) - print("Creating Topics") - print("=" * 50) - - # Create messaging service - print("Creating messaging service...") - sys.stdout.flush() - messaging_service_data = { - "name": "test-messaging-service", - "serviceType": "Kafka", - "connection": { - "config": { - "type": "Kafka", - "bootstrapServers": "localhost:9092" - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/messagingServices", - data=messaging_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - messaging_service_fqn = resp["fullyQualifiedName"] - print(f"Messaging service created: {messaging_service_fqn}") - else: - print(f"Failed to create messaging service: {status} - {resp}") - messaging_service_fqn = None - - if messaging_service_fqn: - # Create topics - print(f"Creating {NUM_TOPICS} topics...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_topic(idx): - topic_data = { - "name": f"topic_{idx:06d}", - "service": messaging_service_fqn, - "partitions": 3, - "replicationFactor": 1, - "description": f"Auto-generated test topic {idx} for distributed indexing testing" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/topics", data=topic_data, method="PUT", headers=headers) - return status in [200, 201] - - # Execute with thread pool - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_topic, i): i for i in range(NUM_TOPICS)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 1000 == 0 or total == NUM_TOPICS: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Topics: {total}/{NUM_TOPICS} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["topics"] = created - print(f"Topics completed: {created} created, {failed} failed") - -# ========== ML MODELS ========== -if NUM_MLMODELS > 0: - print("") - print("=" * 50) - print("Creating ML Models") - print("=" * 50) - - # Create ML model service - print("Creating ML model service...") - sys.stdout.flush() - mlmodel_service_data = { - "name": "test-mlmodel-service", - "serviceType": "Mlflow", - "connection": { - "config": { - "type": "Mlflow", - "trackingUri": "http://mlflow.example.com:5000", - "registryUri": "http://mlflow.example.com:5000" - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/mlmodelServices", - data=mlmodel_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - mlmodel_service_fqn = resp["fullyQualifiedName"] - print(f"ML model service created: {mlmodel_service_fqn}") - else: - print(f"Failed to create ML model service: {status} - {resp}") - mlmodel_service_fqn = None - - if mlmodel_service_fqn: - print(f"Creating {NUM_MLMODELS} ML models...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - algorithms = ["LinearRegression", "RandomForest", "XGBoost", "NeuralNetwork", "SVM", "KMeans", "DecisionTree"] - - def create_mlmodel(idx): - mlmodel_data = { - "name": f"mlmodel_{idx:06d}", - "service": mlmodel_service_fqn, - "algorithm": random.choice(algorithms), - "displayName": f"Test ML Model {idx}", - "description": f"Auto-generated test ML model {idx}" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/mlmodels", data=mlmodel_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_mlmodel, i): i for i in range(NUM_MLMODELS)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 500 == 0 or total == NUM_MLMODELS: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" ML Models: {total}/{NUM_MLMODELS} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["mlmodels"] = created - print(f"ML Models completed: {created} created, {failed} failed") - -# ========== CONTAINERS ========== -if NUM_CONTAINERS > 0: - print("") - print("=" * 50) - print("Creating Containers") - print("=" * 50) - - # Create storage service - print("Creating storage service...") - sys.stdout.flush() - storage_service_data = { - "name": "test-storage-service", - "serviceType": "S3", - "connection": { - "config": { - "type": "S3", - "awsConfig": { - "awsAccessKeyId": "test-key", - "awsSecretAccessKey": "test-secret", - "awsRegion": "us-east-1" - } - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/storageServices", - data=storage_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - storage_service_fqn = resp["fullyQualifiedName"] - print(f"Storage service created: {storage_service_fqn}") - else: - print(f"Failed to create storage service: {status} - {resp}") - storage_service_fqn = None - - if storage_service_fqn: - print(f"Creating {NUM_CONTAINERS} containers...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_container(idx): - container_data = { - "name": f"container_{idx:06d}", - "service": storage_service_fqn, - "displayName": f"Test Container {idx}", - "description": f"Auto-generated test container {idx}" - } - status, _ = make_request(f"{SERVER_URL}/api/v1/containers", data=container_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_container, i): i for i in range(NUM_CONTAINERS)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 500 == 0 or total == NUM_CONTAINERS: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Containers: {total}/{NUM_CONTAINERS} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["containers"] = created - print(f"Containers completed: {created} created, {failed} failed") - -# ========== SEARCH INDEXES ========== -if NUM_SEARCH_INDEXES > 0: - print("") - print("=" * 50) - print("Creating Search Indexes") - print("=" * 50) - - # Create search service - print("Creating search service...") - sys.stdout.flush() - search_service_data = { - "name": "test-search-service", - "serviceType": "ElasticSearch", - "connection": { - "config": { - "type": "ElasticSearch", - "hostPort": "http://elasticsearch.example.com:9200" - } - } - } - - status, resp = make_request( - f"{SERVER_URL}/api/v1/services/searchServices", - data=search_service_data, - method="PUT", - headers=headers - ) - - if status in [200, 201] and isinstance(resp, dict): - search_service_fqn = resp["fullyQualifiedName"] - print(f"Search service created: {search_service_fqn}") - else: - print(f"Failed to create search service: {status} - {resp}") - search_service_fqn = None - - if search_service_fqn: - print(f"Creating {NUM_SEARCH_INDEXES} search indexes...") - sys.stdout.flush() - created = 0 - failed = 0 - start_time = time.time() - - def create_search_index(idx): - search_index_data = { - "name": f"search_index_{idx:06d}", - "service": search_service_fqn, - "displayName": f"Test Search Index {idx}", - "description": f"Auto-generated test search index {idx}", - "fields": [ - {"name": "id", "dataType": "KEYWORD"}, - {"name": "content", "dataType": "TEXT"}, - {"name": "timestamp", "dataType": "DATE"} - ] - } - status, _ = make_request(f"{SERVER_URL}/api/v1/searchIndexes", data=search_index_data, method="PUT", headers=headers) - return status in [200, 201] - - with ThreadPoolExecutor(max_workers=20) as executor: - futures = {executor.submit(create_search_index, i): i for i in range(NUM_SEARCH_INDEXES)} - for future in as_completed(futures): - if future.result(): - created += 1 - else: - failed += 1 - total = created + failed - if total % 200 == 0 or total == NUM_SEARCH_INDEXES: - elapsed = time.time() - start_time - rate = total / elapsed if elapsed > 0 else 0 - print(f" Search Indexes: {total}/{NUM_SEARCH_INDEXES} ({rate:.1f}/sec) - Created: {created}, Failed: {failed}") - sys.stdout.flush() - - stats["searchIndexes"] = created - print(f"Search Indexes completed: {created} created, {failed} failed") - -# ========== SUMMARY ========== -overall_elapsed = time.time() - overall_start -total_created = sum(stats.values()) - -print("") -print("=" * 50) -print("Test Data Loading Complete!") -print("=" * 50) -print("") -print("Summary:") -print(f" Tables: {stats['tables']:>6}") -print(f" Dashboards: {stats['dashboards']:>6}") -print(f" Charts: {stats['charts']:>6}") -print(f" Pipelines: {stats['pipelines']:>6}") -print(f" Topics: {stats['topics']:>6}") -print(f" ML Models: {stats['mlmodels']:>6}") -print(f" Containers: {stats['containers']:>6}") -print(f" Search Indexes: {stats['searchIndexes']:>6}") -print(f" Glossaries: {stats['glossaries']:>6}") -print(f" Glossary Terms: {stats['glossaryTerms']:>6}") -print(f" Classifications: {stats['classifications']:>6}") -print(f" Tags: {stats['tags']:>6}") -print(f" --------------------------") -print(f" Total: {total_created:>6}") -print("") -print(f"Time: {overall_elapsed:.1f} seconds") -if overall_elapsed > 0: - print(f"Rate: {total_created/overall_elapsed:.1f} entities/second") -EOF - -echo "" -echo "Test data loaded. You can now trigger reindexing with: ./scripts/trigger-reindex.sh" diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java index 21aabfa16f4..eadc5ab3267 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java @@ -130,6 +130,7 @@ import org.openmetadata.service.resources.databases.DatasourceConfig; import org.openmetadata.service.resources.filters.ETagRequestFilter; import org.openmetadata.service.resources.filters.ETagResponseFilter; import org.openmetadata.service.resources.settings.SettingsCache; +import org.openmetadata.service.resources.system.DiagnosticsResource; import org.openmetadata.service.search.SearchRepository; import org.openmetadata.service.search.SearchRepositoryFactory; import org.openmetadata.service.secrets.SecretsManagerFactory; @@ -993,6 +994,7 @@ public class OpenMetadataApplication extends Application 0"); + } + if (maxMs < initialMs) { + throw new IllegalArgumentException("maxMs must be >= initialMs"); + } + this.initialMs = initialMs; + this.maxMs = maxMs; + this.currentMs = initialMs; + } + + public long nextDelay() { + long delay = currentMs; + currentMs = Math.min(currentMs * 2, maxMs); + return delay; + } + + public void reset() { + currentMs = initialMs; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreaker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreaker.java new file mode 100644 index 00000000000..32c8b3a90da --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreaker.java @@ -0,0 +1,127 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Iterator; +import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.atomic.AtomicReference; +import lombok.extern.slf4j.Slf4j; + +/** + * Sliding-window circuit breaker for bulk search-index requests. + * + *

State transitions: CLOSED → OPEN (after N failures in window) → HALF_OPEN (probe after + * interval) → CLOSED (on probe success) or back to OPEN (on probe failure). + */ +@Slf4j +public class BulkCircuitBreaker { + + public enum State { + CLOSED, + OPEN, + HALF_OPEN + } + + private final int failureThreshold; + private final long windowMs; + private final long halfOpenProbeMs; + + private final AtomicReference state = new AtomicReference<>(State.CLOSED); + private final ConcurrentLinkedDeque failureTimestamps = new ConcurrentLinkedDeque<>(); + private volatile long openedAt; + + public BulkCircuitBreaker(int failureThreshold, long windowMs, long halfOpenProbeMs) { + if (failureThreshold <= 0) { + throw new IllegalArgumentException("failureThreshold must be > 0"); + } + if (windowMs <= 0) { + throw new IllegalArgumentException("windowMs must be > 0"); + } + if (halfOpenProbeMs <= 0) { + throw new IllegalArgumentException("halfOpenProbeMs must be > 0"); + } + this.failureThreshold = failureThreshold; + this.windowMs = windowMs; + this.halfOpenProbeMs = halfOpenProbeMs; + } + + public boolean allowRequest() { + State current = state.get(); + if (current == State.CLOSED) { + return true; + } + if (current == State.HALF_OPEN) { + return true; + } + // OPEN: check if probe interval has elapsed + if (System.currentTimeMillis() - openedAt >= halfOpenProbeMs) { + if (state.compareAndSet(State.OPEN, State.HALF_OPEN)) { + LOG.warn("Circuit breaker transitioning OPEN → HALF_OPEN (probe request allowed)"); + recordTransition("open_to_half_open"); + } + return true; + } + return false; + } + + public void recordSuccess() { + if (state.compareAndSet(State.HALF_OPEN, State.CLOSED)) { + failureTimestamps.clear(); + LOG.warn("Circuit breaker transitioning HALF_OPEN → CLOSED (probe succeeded)"); + recordTransition("half_open_to_closed"); + } + } + + public void recordFailure() { + long now = System.currentTimeMillis(); + + State current = state.get(); + if (current == State.HALF_OPEN) { + if (state.compareAndSet(State.HALF_OPEN, State.OPEN)) { + openedAt = now; + LOG.warn("Circuit breaker transitioning HALF_OPEN → OPEN (probe failed)"); + recordTransition("half_open_to_open"); + } + return; + } + + failureTimestamps.addLast(now); + pruneOldFailures(now); + + if (failureTimestamps.size() >= failureThreshold + && state.compareAndSet(State.CLOSED, State.OPEN)) { + openedAt = now; + LOG.warn( + "Circuit breaker transitioning CLOSED → OPEN ({} failures in {}ms window)", + failureThreshold, + windowMs); + recordTransition("closed_to_open"); + } + } + + public State getState() { + return state.get(); + } + + public void reset() { + state.set(State.CLOSED); + failureTimestamps.clear(); + } + + private void pruneOldFailures(long now) { + long cutoff = now - windowMs; + Iterator it = failureTimestamps.iterator(); + while (it.hasNext()) { + if (it.next() < cutoff) { + it.remove(); + } else { + break; + } + } + } + + private void recordTransition(String transition) { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordCircuitBreakerTrip(transition); + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkSink.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkSink.java index 144f146acdd..a75274ab98f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkSink.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/BulkSink.java @@ -30,14 +30,20 @@ public interface BulkSink { @FunctionalInterface interface FailureCallback { /** - * Called when a document fails to index in ES/OpenSearch. + * Called when a document fails to index. * * @param entityType The type of entity that failed * @param entityId The ID of the entity (from document ID), may be null for build failures * @param entityFqn The FQN of the entity, may be null if not available - * @param errorMessage The error message from ES/OpenSearch + * @param errorMessage The error message describing the failure + * @param stage The pipeline stage where the failure occurred (PROCESS or SINK) */ - void onFailure(String entityType, String entityId, String entityFqn, String errorMessage); + void onFailure( + String entityType, + String entityId, + String entityFqn, + String errorMessage, + IndexingFailureRecorder.FailureStage stage); } /** @@ -60,6 +66,16 @@ public interface BulkSink { return null; } + /** + * Returns the process stage statistics. This tracks document building/transformation + * separately from the actual sink (bulk indexing) stats. + * + * @return StepStats with process success/failed counts, or null if not supported + */ + default StepStats getProcessStats() { + return null; + } + /** * Wait for all pending vector embedding tasks to complete. This is important for ensuring * no vector tasks are lost when the job completes. The sink's close() method should also @@ -82,6 +98,30 @@ public interface BulkSink { return 0; } + /** + * Returns the number of currently active (in-flight) bulk requests. + * + * @return Number of active bulk requests + */ + default int getActiveBulkRequestCount() { + return 0; + } + + /** + * Wait for vector embedding tasks to complete and return detailed result including timing. + * + * @param timeoutSeconds Maximum time to wait + * @return VectorCompletionResult with completion status, pending count, and wait time + */ + default VectorCompletionResult awaitVectorCompletionWithDetails(int timeoutSeconds) { + long start = System.currentTimeMillis(); + boolean ok = awaitVectorCompletion(timeoutSeconds); + long waited = System.currentTimeMillis() - start; + return ok + ? VectorCompletionResult.success(waited) + : VectorCompletionResult.timeout(getPendingVectorTaskCount(), waited); + } + /** Key for passing StageStatsTracker through context data to the sink. */ String STATS_TRACKER_CONTEXT_KEY = "stageStatsTracker"; } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListener.java index a373df78305..d2614122d81 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListener.java @@ -122,6 +122,51 @@ public class CompositeProgressListener implements ReindexingProgressListener { } } + @Override + public void onReaderFailure(String entityType, String entityId, String error, FailureType type) { + for (ReindexingProgressListener listener : listeners) { + try { + listener.onReaderFailure(entityType, entityId, error, type); + } catch (Exception e) { + LOG.error("Listener {} failed on onReaderFailure", listener.getClass().getSimpleName(), e); + } + } + } + + @Override + public void onProcessFailure(String entityType, String entityId, String error) { + for (ReindexingProgressListener listener : listeners) { + try { + listener.onProcessFailure(entityType, entityId, error); + } catch (Exception e) { + LOG.error("Listener {} failed on onProcessFailure", listener.getClass().getSimpleName(), e); + } + } + } + + @Override + public void onSinkFailure(String entityType, String entityId, String error) { + for (ReindexingProgressListener listener : listeners) { + try { + listener.onSinkFailure(entityType, entityId, error); + } catch (Exception e) { + LOG.error("Listener {} failed on onSinkFailure", listener.getClass().getSimpleName(), e); + } + } + } + + @Override + public void onSubIndexingCompleted(String entityType, String subIndex, StepStats subIndexStats) { + for (ReindexingProgressListener listener : listeners) { + try { + listener.onSubIndexingCompleted(entityType, subIndex, subIndexStats); + } catch (Exception e) { + LOG.error( + "Listener {} failed on onSubIndexingCompleted", listener.getClass().getSimpleName(), e); + } + } + } + @Override public void onJobCompleted(Stats finalStats, long elapsedMillis) { for (ReindexingProgressListener listener : listeners) { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java new file mode 100644 index 00000000000..fabe125e4e6 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java @@ -0,0 +1,649 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; + +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.type.Include; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.DistributedSearchIndexExecutor; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexJob; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.search.vector.VectorIndexService; +import org.openmetadata.service.util.FullyQualifiedName; + +@Slf4j +public class DistributedIndexingStrategy implements IndexingStrategy { + + private static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + + private static final long MONITOR_POLL_INTERVAL_MS = 2000; + + private final CollectionDAO collectionDAO; + private final SearchRepository searchRepository; + private final EventPublisherJob jobData; + private final UUID appId; + private final Long appStartTime; + private final String createdBy; + + private final CompositeProgressListener listeners = new CompositeProgressListener(); + private final AtomicBoolean stopped = new AtomicBoolean(false); + private final AtomicReference currentStats = new AtomicReference<>(); + + private volatile DistributedSearchIndexExecutor distributedExecutor; + private volatile BulkSink searchIndexSink; + private volatile ReindexingConfiguration config; + + public DistributedIndexingStrategy( + CollectionDAO collectionDAO, + SearchRepository searchRepository, + EventPublisherJob jobData, + UUID appId, + Long appStartTime, + String createdBy) { + this.collectionDAO = collectionDAO; + this.searchRepository = searchRepository; + this.jobData = jobData; + this.appId = appId; + this.appStartTime = appStartTime; + this.createdBy = createdBy; + } + + @Override + public void addListener(ReindexingProgressListener listener) { + listeners.addListener(listener); + } + + @Override + public ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context) { + long startTime = System.currentTimeMillis(); + try { + return doExecute(config, context, startTime); + } catch (Exception e) { + LOG.error("Distributed reindexing failed", e); + Stats stats = currentStats.get(); + return ExecutionResult.fromStats(stats, ExecutionResult.Status.FAILED, startTime); + } + } + + private ExecutionResult doExecute( + ReindexingConfiguration config, ReindexingJobContext context, long startTime) + throws Exception { + + this.config = config; + LOG.info("Starting distributed reindexing for entities: {}", config.entities()); + + Stats stats = initializeTotalRecords(config.entities()); + currentStats.set(stats); + + int partitionSize = jobData.getPartitionSize() != null ? jobData.getPartitionSize() : 10000; + distributedExecutor = new DistributedSearchIndexExecutor(collectionDAO, partitionSize); + distributedExecutor.performStartupRecovery(); + + distributedExecutor.addListener(listeners); + + SearchIndexJob distributedJob = + distributedExecutor.createJob(config.entities(), jobData, createdBy, config); + + LOG.info( + "Created distributed job {} with {} total records", + distributedJob.getId(), + distributedJob.getTotalRecords()); + + searchIndexSink = + searchRepository.createBulkSink( + config.batchSize(), config.maxConcurrentRequests(), config.payloadSize()); + + RecreateIndexHandler recreateIndexHandler = searchRepository.createReindexHandler(); + ReindexContext recreateContext = null; + + if (config.recreateIndex()) { + recreateContext = recreateIndexHandler.reCreateIndexes(config.entities()); + if (recreateContext != null && !recreateContext.isEmpty()) { + distributedExecutor.updateStagedIndexMapping(recreateContext.getStagedIndexMapping()); + } + } + + distributedExecutor.setAppContext(appId, appStartTime); + distributedExecutor.execute( + searchIndexSink, recreateContext, Boolean.TRUE.equals(config.recreateIndex()), config); + + monitorDistributedJob(distributedJob.getId()); + + flushAndAwaitSink(); + + SearchIndexJob finalJob = distributedExecutor.getJobWithFreshStats(); + Map metadata = new HashMap<>(); + + if (finalJob != null) { + StepStats sinkStats = searchIndexSink != null ? searchIndexSink.getStats() : null; + updateStatsFromDistributedJob(stats, finalJob, sinkStats); + + if (searchIndexSink != null) { + StepStats sinkVectorStats = searchIndexSink.getVectorStats(); + if (sinkVectorStats != null && sinkVectorStats.getTotalRecords() > 0) { + stats.setVectorStats(sinkVectorStats); + } + } + + if (finalJob.getServerStats() != null && !finalJob.getServerStats().isEmpty()) { + metadata.put("serverStats", finalJob.getServerStats()); + metadata.put("serverCount", finalJob.getServerStats().size()); + metadata.put("distributedJobId", finalJob.getId().toString()); + } + } + + currentStats.set(stats); + + boolean success = + finalizeAllEntityReindex( + recreateIndexHandler, + recreateContext, + !stopped.get() && !hasIncompleteProcessing(stats)); + + ExecutionResult.Status resultStatus = determineStatus(stats); + + StatsReconciler.reconcile(stats); + + return ExecutionResult.builder() + .status(resultStatus) + .totalRecords(stats.getJobStats().getTotalRecords()) + .successRecords(stats.getJobStats().getSuccessRecords()) + .failedRecords(stats.getJobStats().getFailedRecords()) + .startTime(startTime) + .endTime(System.currentTimeMillis()) + .finalStats(stats) + .metadata(metadata) + .build(); + } + + private void flushAndAwaitSink() { + if (searchIndexSink == null) { + return; + } + + int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); + if (pendingVectorTasks > 0) { + LOG.info("Waiting for {} pending vector embedding tasks to complete", pendingVectorTasks); + boolean vectorComplete = searchIndexSink.awaitVectorCompletion(120); + if (!vectorComplete) { + LOG.warn("Vector embedding wait timed out - some tasks may not be reflected in stats"); + } + } + + LOG.info("Flushing sink and waiting for pending bulk requests"); + boolean flushComplete = searchIndexSink.flushAndAwait(60); + if (!flushComplete) { + LOG.warn("Sink flush timed out - some requests may not be reflected in stats"); + } + + try { + searchIndexSink.close(); + } catch (Exception e) { + LOG.error("Error closing search index sink", e); + } + } + + private void monitorDistributedJob(UUID jobId) { + CountDownLatch completionLatch = new CountDownLatch(1); + ScheduledExecutorService monitor = + Executors.newSingleThreadScheduledExecutor( + Thread.ofPlatform().name("distributed-monitor").factory()); + + try { + monitor.scheduleAtFixedRate( + () -> { + try { + if (stopped.get()) { + LOG.info("Stop signal received, stopping distributed job"); + distributedExecutor.stop(); + completionLatch.countDown(); + return; + } + + SearchIndexJob job = distributedExecutor.getJobWithFreshStats(); + if (job == null) { + completionLatch.countDown(); + return; + } + + IndexJobStatus status = job.getStatus(); + if (status == IndexJobStatus.COMPLETED + || status == IndexJobStatus.COMPLETED_WITH_ERRORS + || status == IndexJobStatus.FAILED + || status == IndexJobStatus.STOPPED) { + LOG.info("Distributed job {} completed with status: {}", jobId, status); + completionLatch.countDown(); + return; + } + + updateStatsFromDistributedJob(currentStats.get(), job, null); + } catch (Exception e) { + LOG.error("Error in distributed job monitor task for job {}", jobId, e); + } + }, + 0, + MONITOR_POLL_INTERVAL_MS, + TimeUnit.MILLISECONDS); + + completionLatch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warn("Distributed job monitoring interrupted"); + } finally { + monitor.shutdownNow(); + } + } + + private void updateStatsFromDistributedJob( + Stats stats, SearchIndexJob distributedJob, StepStats actualSinkStats) { + if (stats == null) { + return; + } + + CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats serverStatsAggr = null; + try { + serverStatsAggr = + Entity.getCollectionDAO() + .searchIndexServerStatsDAO() + .getAggregatedStats(distributedJob.getId().toString()); + if (serverStatsAggr != null) { + LOG.info( + "Fetched aggregated server stats for job {}: readerSuccess={}, readerFailed={}, " + + "sinkSuccess={}, sinkFailed={}", + distributedJob.getId(), + serverStatsAggr.readerSuccess(), + serverStatsAggr.readerFailed(), + serverStatsAggr.sinkSuccess(), + serverStatsAggr.sinkFailed()); + } + } catch (Exception e) { + LOG.debug("Could not fetch aggregated server stats for job {}", distributedJob.getId(), e); + } + + long successRecords; + long failedRecords; + String statsSource; + + if (serverStatsAggr != null && serverStatsAggr.sinkSuccess() > 0) { + successRecords = serverStatsAggr.sinkSuccess(); + failedRecords = + serverStatsAggr.readerFailed() + + serverStatsAggr.sinkFailed() + + serverStatsAggr.processFailed(); + statsSource = "serverStatsTable"; + } else if (actualSinkStats != null) { + successRecords = actualSinkStats.getSuccessRecords(); + failedRecords = actualSinkStats.getFailedRecords(); + statsSource = "localSink"; + } else { + successRecords = distributedJob.getSuccessRecords(); + failedRecords = distributedJob.getFailedRecords(); + statsSource = "partition-based"; + } + + LOG.debug( + "Stats source: {}, success={}, failed={}", statsSource, successRecords, failedRecords); + + StepStats jobStats = stats.getJobStats(); + if (jobStats != null) { + jobStats.setSuccessRecords(saturatedToInt(successRecords)); + jobStats.setFailedRecords(saturatedToInt(failedRecords)); + } + + StepStats readerStats = stats.getReaderStats(); + if (readerStats != null) { + readerStats.setTotalRecords(saturatedToInt(distributedJob.getTotalRecords())); + long readerFailed = serverStatsAggr != null ? serverStatsAggr.readerFailed() : 0; + long readerWarnings = serverStatsAggr != null ? serverStatsAggr.readerWarnings() : 0; + long readerSuccess = + serverStatsAggr != null + ? serverStatsAggr.readerSuccess() + : distributedJob.getTotalRecords() - readerFailed - readerWarnings; + readerStats.setSuccessRecords(saturatedToInt(readerSuccess)); + readerStats.setFailedRecords(saturatedToInt(readerFailed)); + readerStats.setWarningRecords(saturatedToInt(readerWarnings)); + } + + StepStats processStats = stats.getProcessStats(); + if (processStats != null && serverStatsAggr != null) { + long processSuccess = serverStatsAggr.processSuccess(); + long processFailed = serverStatsAggr.processFailed(); + processStats.setTotalRecords(saturatedToInt(processSuccess + processFailed)); + processStats.setSuccessRecords(saturatedToInt(processSuccess)); + processStats.setFailedRecords(saturatedToInt(processFailed)); + } + + StepStats sinkStats = stats.getSinkStats(); + if (sinkStats != null) { + if (serverStatsAggr != null) { + long sinkSuccess = serverStatsAggr.sinkSuccess(); + long sinkFailed = serverStatsAggr.sinkFailed(); + long actualSinkTotal = sinkSuccess + sinkFailed; + sinkStats.setTotalRecords(saturatedToInt(actualSinkTotal)); + sinkStats.setSuccessRecords(saturatedToInt(sinkSuccess)); + sinkStats.setFailedRecords(saturatedToInt(sinkFailed)); + } else { + long sinkTotal = distributedJob.getTotalRecords(); + sinkStats.setTotalRecords(saturatedToInt(sinkTotal)); + sinkStats.setSuccessRecords(saturatedToInt(successRecords)); + sinkStats.setFailedRecords(saturatedToInt(failedRecords)); + } + } + + StepStats vectorStats = stats.getVectorStats(); + if (vectorStats != null && serverStatsAggr != null) { + long vectorSuccess = serverStatsAggr.vectorSuccess(); + long vectorFailed = serverStatsAggr.vectorFailed(); + vectorStats.setTotalRecords(saturatedToInt(vectorSuccess + vectorFailed)); + vectorStats.setSuccessRecords(saturatedToInt(vectorSuccess)); + vectorStats.setFailedRecords(saturatedToInt(vectorFailed)); + } + + if (distributedJob.getEntityStats() != null && stats.getEntityStats() != null) { + for (Map.Entry entry : + distributedJob.getEntityStats().entrySet()) { + StepStats entityStats = + stats.getEntityStats().getAdditionalProperties().get(entry.getKey()); + if (entityStats != null) { + entityStats.setSuccessRecords(saturatedToInt(entry.getValue().getSuccessRecords())); + entityStats.setFailedRecords(saturatedToInt(entry.getValue().getFailedRecords())); + } + } + } + + StatsReconciler.reconcile(stats); + } + + private static int saturatedToInt(long value) { + return (int) Math.min(value, Integer.MAX_VALUE); + } + + private ExecutionResult.Status determineStatus(Stats stats) { + if (stopped.get()) { + return ExecutionResult.Status.STOPPED; + } + if (hasIncompleteProcessing(stats)) { + return ExecutionResult.Status.COMPLETED_WITH_ERRORS; + } + return ExecutionResult.Status.COMPLETED; + } + + private boolean hasIncompleteProcessing(Stats stats) { + if (stats == null || stats.getJobStats() == null) { + return false; + } + StepStats jobStats = stats.getJobStats(); + long failed = jobStats.getFailedRecords() != null ? jobStats.getFailedRecords() : 0; + long processed = jobStats.getSuccessRecords() != null ? jobStats.getSuccessRecords() : 0; + long total = jobStats.getTotalRecords() != null ? jobStats.getTotalRecords() : 0; + return failed > 0 || (total > 0 && processed < total); + } + + private boolean finalizeAllEntityReindex( + RecreateIndexHandler recreateIndexHandler, + ReindexContext recreateContext, + boolean finalSuccess) { + if (recreateIndexHandler == null || recreateContext == null) { + return finalSuccess; + } + + Set promotedEntities = Collections.emptySet(); + if (distributedExecutor != null && distributedExecutor.getEntityTracker() != null) { + promotedEntities = distributedExecutor.getEntityTracker().getPromotedEntities(); + } + + // Get per-entity stats for determining per-entity success + Map entityStatsMap = Collections.emptyMap(); + if (distributedExecutor != null) { + SearchIndexJob finalJob = distributedExecutor.getJobWithFreshStats(); + if (finalJob != null && finalJob.getEntityStats() != null) { + entityStatsMap = finalJob.getEntityStats(); + } + } + + LOG.debug( + "Finalization: finalSuccess={}, promotedEntities={}, allEntities={}", + finalSuccess, + promotedEntities, + recreateContext.getEntities()); + + Set entitiesToFinalize = new HashSet<>(recreateContext.getEntities()); + entitiesToFinalize.removeAll(promotedEntities); + + boolean hasVectorIndex = entitiesToFinalize.remove(VectorIndexService.VECTOR_INDEX_KEY); + + LOG.debug("Entities to finalize={}, already promoted={}", entitiesToFinalize, promotedEntities); + + try { + if (!entitiesToFinalize.isEmpty()) { + LOG.info( + "Finalizing {} remaining entities (already promoted: {})", + entitiesToFinalize.size(), + promotedEntities.size()); + + for (String entityType : entitiesToFinalize) { + try { + boolean entitySuccess = computeEntitySuccess(entityType, entityStatsMap); + LOG.debug( + "Finalizing entity '{}' with perEntitySuccess={} (globalSuccess={})", + entityType, + entitySuccess, + finalSuccess); + finalizeEntityReindex(recreateIndexHandler, recreateContext, entityType, entitySuccess); + } catch (Exception ex) { + LOG.error("Failed to finalize reindex for entity: {}", entityType, ex); + } + } + } + + if (hasVectorIndex) { + boolean vectorSuccess = + finalSuccess + || (currentStats.get() != null && !hasIncompleteProcessing(currentStats.get())); + try { + finalizeEntityReindex( + recreateIndexHandler, + recreateContext, + VectorIndexService.VECTOR_INDEX_KEY, + vectorSuccess); + } catch (Exception ex) { + LOG.error("Failed to finalize vector index", ex); + } + } + } catch (Exception e) { + LOG.error("Error during entity finalization", e); + } + + return finalSuccess; + } + + private boolean computeEntitySuccess( + String entityType, Map entityStatsMap) { + if (entityStatsMap == null || entityStatsMap.isEmpty()) { + return false; + } + SearchIndexJob.EntityTypeStats stats = entityStatsMap.get(entityType); + if (stats == null) { + // Entity not in stats means 0 records — nothing to index = success + return true; + } + return stats.getFailedRecords() == 0 + && stats.getSuccessRecords() + stats.getFailedRecords() >= stats.getTotalRecords(); + } + + private void finalizeEntityReindex( + RecreateIndexHandler recreateIndexHandler, + ReindexContext recreateContext, + String entityType, + boolean success) { + try { + var entityReindexContext = + org.openmetadata.service.search.EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) + .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases( + new HashSet<>(listOrEmpty(recreateContext.getParentAliases(entityType)))) + .build(); + + recreateIndexHandler.finalizeReindex(entityReindexContext, success); + } catch (Exception ex) { + LOG.error("Failed to finalize index recreation flow for {}", entityType, ex); + } + } + + @Override + public Optional getStats() { + return Optional.ofNullable(currentStats.get()); + } + + @Override + public void stop() { + if (stopped.compareAndSet(false, true)) { + LOG.info("Stopping distributed indexing strategy"); + + if (distributedExecutor != null) { + try { + distributedExecutor.stop(); + } catch (Exception e) { + LOG.error("Error stopping distributed executor", e); + } + } + // Do NOT close the sink here — workers may still be writing to it. + // The sink is properly flushed and closed by flushAndAwaitSink() in doExecute() + // after the monitor exits and the executor's finally block completes. + } + } + + @Override + public boolean isStopped() { + return stopped.get(); + } + + Stats initializeTotalRecords(Set entities) { + Stats stats = new Stats(); + stats.setEntityStats(new org.openmetadata.schema.system.EntityStats()); + stats.setJobStats(new StepStats()); + stats.setReaderStats(new StepStats()); + stats.setProcessStats(new StepStats()); + stats.setSinkStats(new StepStats()); + stats.setVectorStats(new StepStats()); + + List ordered = EntityPriority.sortByPriority(entities); + int total = 0; + for (String entityType : ordered) { + int entityTotal = getEntityTotal(entityType); + total += entityTotal; + + StepStats entityStats = new StepStats(); + entityStats.setTotalRecords(entityTotal); + entityStats.setSuccessRecords(0); + entityStats.setFailedRecords(0); + stats.getEntityStats().getAdditionalProperties().put(entityType, entityStats); + } + + stats.getJobStats().setTotalRecords(total); + stats.getJobStats().setSuccessRecords(0); + stats.getJobStats().setFailedRecords(0); + + stats.getReaderStats().setTotalRecords(total); + stats.getReaderStats().setSuccessRecords(0); + stats.getReaderStats().setFailedRecords(0); + + stats.getProcessStats().setTotalRecords(0); + stats.getProcessStats().setSuccessRecords(0); + stats.getProcessStats().setFailedRecords(0); + + stats.getSinkStats().setTotalRecords(0); + stats.getSinkStats().setSuccessRecords(0); + stats.getSinkStats().setFailedRecords(0); + + stats.getVectorStats().setTotalRecords(0); + stats.getVectorStats().setSuccessRecords(0); + stats.getVectorStats().setFailedRecords(0); + + return stats; + } + + private int getEntityTotal(String entityType) { + try { + String correctedType = "queryCostResult".equals(entityType) ? QUERY_COST_RECORD : entityType; + + if (!TIME_SERIES_ENTITIES.contains(correctedType)) { + return Entity.getEntityRepository(correctedType) + .getDao() + .listCount(new ListFilter(Include.ALL)); + } else { + ListFilter listFilter = new ListFilter(null); + EntityTimeSeriesRepository repository; + + if (isDataInsightIndex(correctedType)) { + listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(correctedType)); + repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); + } else { + repository = Entity.getEntityTimeSeriesRepository(correctedType); + } + + if (config != null) { + long startTs = config.getTimeSeriesStartTs(correctedType); + if (startTs > 0) { + long endTs = System.currentTimeMillis(); + return repository.getTimeSeriesDao().listCount(listFilter, startTs, endTs, false); + } + } + return repository.getTimeSeriesDao().listCount(listFilter); + } + } catch (Exception e) { + LOG.debug("Error getting total for '{}'", entityType, e); + return 0; + } + } + + private boolean isDataInsightIndex(String entityType) { + return entityType.endsWith("ReportData"); + } + + DistributedSearchIndexExecutor getDistributedExecutor() { + return distributedExecutor; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java index 353d703daac..13bd1d6e9b7 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java @@ -66,6 +66,10 @@ public class ElasticSearchBulkSink implements BulkSink { private final AtomicLong totalSuccess = new AtomicLong(0); private final AtomicLong totalFailed = new AtomicLong(0); + // Process stage metrics (document building/transformation) + private final AtomicLong processSuccess = new AtomicLong(0); + private final AtomicLong processFailed = new AtomicLong(0); + // Configuration private volatile int batchSize; private volatile int maxConcurrentRequests; @@ -99,6 +103,7 @@ public class ElasticSearchBulkSink implements BulkSink { concurrentRequests, maxPayloadSizeBytes / (1024 * 1024)); + BulkCircuitBreaker circuitBreaker = new BulkCircuitBreaker(5, 30_000, 10_000); return new CustomBulkProcessor( searchClient, bulkActions, @@ -110,7 +115,8 @@ public class ElasticSearchBulkSink implements BulkSink { totalSubmitted, totalSuccess, totalFailed, - this::updateStats); + this::updateStats, + circuitBreaker); } @Override @@ -258,12 +264,14 @@ public class ElasticSearchBulkSink implements BulkSink { tracker.incrementPendingSink(); } bulkProcessor.add(operation, docId, entityType, tracker, estimatedSize); + processSuccess.incrementAndGet(); if (tracker != null) { tracker.recordProcess(StatsResult.SUCCESS); } } catch (EntityNotFoundException e) { LOG.error("Entity Not Found Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -274,12 +282,14 @@ public class ElasticSearchBulkSink implements BulkSink { entityTypeName, entity.getId() != null ? entity.getId().toString() : null, entity.getFullyQualifiedName(), - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } catch (Exception e) { LOG.error( "Encountered Issue while building SearchDoc from Entity Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -290,7 +300,8 @@ public class ElasticSearchBulkSink implements BulkSink { entityTypeName, entity.getId() != null ? entity.getId().toString() : null, entity.getFullyQualifiedName(), - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } } @@ -317,12 +328,14 @@ public class ElasticSearchBulkSink implements BulkSink { tracker.incrementPendingSink(); } bulkProcessor.add(operation, docId, entityType, tracker, estimatedSize); + processSuccess.incrementAndGet(); if (tracker != null) { tracker.recordProcess(StatsResult.SUCCESS); } } catch (EntityNotFoundException e) { LOG.error("Entity Not Found Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -332,12 +345,14 @@ public class ElasticSearchBulkSink implements BulkSink { entityType, entity.getId() != null ? entity.getId().toString() : null, null, - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } catch (Exception e) { LOG.error( "Encountered Issue while building SearchDoc from Entity Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -347,7 +362,8 @@ public class ElasticSearchBulkSink implements BulkSink { entityType, entity.getId() != null ? entity.getId().toString() : null, null, - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } } @@ -377,6 +393,16 @@ public class ElasticSearchBulkSink implements BulkSink { .withFailedRecords((int) failed); } + @Override + public StepStats getProcessStats() { + long success = processSuccess.get(); + long failed = processFailed.get(); + return new StepStats() + .withTotalRecords((int) (success + failed)) + .withSuccessRecords((int) success) + .withFailedRecords((int) failed); + } + @Override public void close() { try { @@ -404,6 +430,11 @@ public class ElasticSearchBulkSink implements BulkSink { } } + @Override + public int getActiveBulkRequestCount() { + return bulkProcessor.activeBulkRequests.get(); + } + @Override public boolean flushAndAwait(int timeoutSeconds) { try { @@ -501,6 +532,7 @@ public class ElasticSearchBulkSink implements BulkSink { private final int maxRetries; private volatile boolean closed = false; private volatile FailureCallback failureCallback; + private final BulkCircuitBreaker circuitBreaker; CustomBulkProcessor( ElasticSearchClient client, @@ -513,7 +545,8 @@ public class ElasticSearchBulkSink implements BulkSink { AtomicLong totalSubmitted, AtomicLong totalSuccess, AtomicLong totalFailed, - Runnable statsUpdater) { + Runnable statsUpdater, + BulkCircuitBreaker circuitBreaker) { this.asyncClient = new ElasticsearchAsyncClient(client.getNewClient()._transport()); this.bulkActions = bulkActions; this.maxPayloadSizeBytes = maxPayloadSizeBytes; @@ -524,6 +557,7 @@ public class ElasticSearchBulkSink implements BulkSink { this.totalSuccess = totalSuccess; this.totalFailed = totalFailed; this.statsUpdater = statsUpdater; + this.circuitBreaker = circuitBreaker; this.scheduler = Executors.newScheduledThreadPool(1); scheduler.scheduleAtFixedRate( @@ -668,17 +702,35 @@ public class ElasticSearchBulkSink implements BulkSink { private void executeBulkWithRetry( List operations, long executionId, int numberOfActions, int attemptNumber) { + if (!circuitBreaker.allowRequest()) { + LOG.warn( + "Circuit breaker OPEN - fail-fast for bulk request {} with {} actions", + executionId, + numberOfActions); + totalFailed.addAndGet(numberOfActions); + statsUpdater.run(); + activeBulkRequests.decrementAndGet(); + concurrentRequestSemaphore.release(); + return; + } + CompletableFuture future = asyncClient.bulk(b -> b.operations(operations).refresh(Refresh.False)); future.whenComplete( (response, error) -> { + boolean retryScheduled = false; try { if (error != null) { - handleBulkFailure(operations, executionId, numberOfActions, attemptNumber, error); + circuitBreaker.recordFailure(); + retryScheduled = + handleBulkFailure( + operations, executionId, numberOfActions, attemptNumber, error); } else if (response.errors()) { + circuitBreaker.recordSuccess(); handlePartialFailure(response, executionId, numberOfActions); } else { + circuitBreaker.recordSuccess(); totalSuccess.addAndGet(numberOfActions); LOG.debug( "Bulk request {} completed successfully with {} actions", @@ -707,9 +759,7 @@ public class ElasticSearchBulkSink implements BulkSink { } } } finally { - if (error != null && shouldRetry(attemptNumber, error)) { - // Don't release resources yet, we're retrying - } else { + if (!retryScheduled) { activeBulkRequests.decrementAndGet(); concurrentRequestSemaphore.release(); } @@ -717,13 +767,14 @@ public class ElasticSearchBulkSink implements BulkSink { }); } - private void handleBulkFailure( + private boolean handleBulkFailure( List operations, long executionId, int numberOfActions, int attemptNumber, Throwable error) { - if (shouldRetry(attemptNumber, error)) { + if (shouldRetry(attemptNumber, error) + && circuitBreaker.getState() != BulkCircuitBreaker.State.OPEN) { long backoffTime = calculateBackoff(attemptNumber); LOG.warn( "Bulk request {} failed (attempt {}), retrying in {}ms: {}", @@ -736,6 +787,7 @@ public class ElasticSearchBulkSink implements BulkSink { () -> executeBulkWithRetry(operations, executionId, numberOfActions, attemptNumber + 1), backoffTime, TimeUnit.MILLISECONDS); + return true; } else { totalFailed.addAndGet(numberOfActions); LOG.error( @@ -758,11 +810,17 @@ public class ElasticSearchBulkSink implements BulkSink { tracker.recordSink(StatsResult.FAILED); } if (failureCallback != null) { - failureCallback.onFailure(entityType, docId, null, error.getMessage()); + failureCallback.onFailure( + entityType, + docId, + null, + error.getMessage(), + IndexingFailureRecorder.FailureStage.SINK); } } } statsUpdater.run(); + return false; } } @@ -791,7 +849,8 @@ public class ElasticSearchBulkSink implements BulkSink { tracker.recordSink(StatsResult.FAILED); } if (failureCallback != null) { - failureCallback.onFailure(entityType, docId, null, failureMessage); + failureCallback.onFailure( + entityType, docId, null, failureMessage, IndexingFailureRecorder.FailureStage.SINK); } } else { // Clean up on success diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java new file mode 100644 index 00000000000..512d990ee2a --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java @@ -0,0 +1,38 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Set; + +/** + * Per-entity-type batch sizing based on typical document size. Large entity types (tables, + * dashboards, etc.) produce bigger search documents, so we use smaller batches. Small entity types + * (users, tags, etc.) produce tiny documents, so we can use larger batches. + */ +public final class EntityBatchSizeEstimator { + + private static final Set LARGE_ENTITIES = + Set.of("table", "topic", "dashboard", "mlmodel", "container", "storedProcedure"); + + private static final Set SMALL_ENTITIES = + Set.of("user", "team", "bot", "role", "policy", "tag", "classification"); + + private static final int MIN_BATCH_SIZE = 25; + private static final int MAX_BATCH_SIZE = 1000; + + private EntityBatchSizeEstimator() {} + + public static int estimateBatchSize(String entityType, int baseBatchSize) { + if (baseBatchSize <= 0) { + return baseBatchSize; + } + + if (LARGE_ENTITIES.contains(entityType)) { + return Math.max(baseBatchSize / 2, MIN_BATCH_SIZE); + } + + if (SMALL_ENTITIES.contains(entityType)) { + return Math.min(baseBatchSize * 2, MAX_BATCH_SIZE); + } + + return baseBatchSize; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriority.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriority.java new file mode 100644 index 00000000000..38cfda9aed2 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriority.java @@ -0,0 +1,125 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Defines priority tiers for entity types during reindexing. Higher-priority entities are processed + * first so that foundational data (services, users) is available in the search index before + * dependent data (tables, dashboards) is indexed. + */ +public final class EntityPriority { + + enum Tier { + CRITICAL(0, 100), + HIGH(1, 80), + MEDIUM(2, 60), + LOW(3, 40), + LOWEST(4, 20); + + private final int order; + private final int numericPriority; + + Tier(int order, int numericPriority) { + this.order = order; + this.numericPriority = numericPriority; + } + + int order() { + return order; + } + + int numericPriority() { + return numericPriority; + } + } + + private static final Map ENTITY_TIERS = new LinkedHashMap<>(); + + static { + // P0 CRITICAL: Service entities — hierarchy parents, must exist first + for (String s : + List.of( + "databaseService", + "messagingService", + "dashboardService", + "pipelineService", + "mlmodelService", + "storageService", + "searchService", + "apiService", + "metadataService")) { + ENTITY_TIERS.put(s, Tier.CRITICAL); + } + + // P1 HIGH: Identity/org entities — referenced by everything + for (String s : List.of("user", "team", "role", "bot", "persona")) { + ENTITY_TIERS.put(s, Tier.HIGH); + } + + // P2 MEDIUM: Core data assets + for (String s : + List.of( + "table", + "database", + "databaseSchema", + "dashboard", + "chart", + "pipeline", + "topic", + "mlmodel", + "container", + "storedProcedure", + "query", + "dashboardDataModel", + "api", + "apiEndpoint", + "apiCollection")) { + ENTITY_TIERS.put(s, Tier.MEDIUM); + } + + // P4 LOWEST: Time series entities + for (String s : + List.of( + "entityReportData", + "rawCostAnalysisReportData", + "webAnalyticUserActivityReportData", + "webAnalyticEntityViewReportData", + "aggregatedCostAnalysisReportData", + "testCaseResolutionStatus", + "testCaseResult", + "queryCostRecord")) { + ENTITY_TIERS.put(s, Tier.LOWEST); + } + + // P3 LOW: Everything else not explicitly listed defaults to LOW via getTier() + } + + private EntityPriority() {} + + static Tier getTier(String entityType) { + return ENTITY_TIERS.getOrDefault(entityType, Tier.LOW); + } + + /** + * Returns a numeric priority score for the given entity type. Higher values mean higher priority. + * Used by the distributed indexing path to store priority on partitions for SQL-level ordering. + */ + public static int getNumericPriority(String entityType) { + return getTier(entityType).numericPriority(); + } + + /** + * Returns entity types sorted by priority tier. Entities within the same tier preserve their + * original iteration order from the input set. + */ + public static List sortByPriority(Set entities) { + List list = new ArrayList<>(entities); + list.sort(Comparator.comparingInt(e -> getTier(e).order())); + return list; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java new file mode 100644 index 00000000000..c77b78d5b22 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java @@ -0,0 +1,336 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Phaser; +import java.util.concurrent.atomic.AtomicBoolean; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.util.RestUtil; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; + +/** + * Standalone reader that encapsulates all entity reading logic. Decoupled from queues and sinks — + * delivers batches via a callback interface. + */ +@Slf4j +public class EntityReader implements AutoCloseable { + + static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + + private static final int MAX_READERS_PER_ENTITY = 5; + + @FunctionalInterface + public interface BatchCallback { + void onBatchRead(String entityType, ResultList batch, int offset) + throws InterruptedException; + } + + @FunctionalInterface + interface KeysetBatchReader { + ResultList readNextKeyset(String cursor) throws SearchIndexException; + } + + @FunctionalInterface + interface BoundaryFinder { + List findBoundaries(int numReaders, int totalRecords); + } + + private static final int DEFAULT_MAX_RETRY_ATTEMPTS = 3; + private static final long DEFAULT_RETRY_BACKOFF_MS = 500; + + private final ExecutorService producerExecutor; + private final AtomicBoolean stopped; + private final int maxRetryAttempts; + private final long retryBackoffMs; + + public EntityReader(ExecutorService producerExecutor, AtomicBoolean stopped) { + this(producerExecutor, stopped, DEFAULT_MAX_RETRY_ATTEMPTS, DEFAULT_RETRY_BACKOFF_MS); + } + + public EntityReader( + ExecutorService producerExecutor, + AtomicBoolean stopped, + int maxRetryAttempts, + long retryBackoffMs) { + this.producerExecutor = producerExecutor; + this.stopped = stopped; + this.maxRetryAttempts = maxRetryAttempts; + this.retryBackoffMs = retryBackoffMs; + } + + /** + * Read all entities of a given type, invoking callback for each batch. + * + * @param entityType The entity type to read + * @param totalRecords Total records expected for this entity + * @param batchSize Batch size for reading + * @param phaser Phaser for completion tracking (readers will register/deregister) + * @param callback Callback invoked with each batch + * @return Number of readers submitted + */ + public int readEntity( + String entityType, int totalRecords, int batchSize, Phaser phaser, BatchCallback callback) { + return readEntity(entityType, totalRecords, batchSize, phaser, callback, null, null); + } + + public int readEntity( + String entityType, + int totalRecords, + int batchSize, + Phaser phaser, + BatchCallback callback, + Long timeSeriesStartTs, + Long timeSeriesEndTs) { + if (totalRecords <= 0) { + return 0; + } + + int numReaders = + Math.min(calculateNumberOfReaders(totalRecords, batchSize), MAX_READERS_PER_ENTITY); + phaser.bulkRegister(numReaders); + + try { + if (TIME_SERIES_ENTITIES.contains(entityType)) { + submitReaders( + entityType, + totalRecords, + batchSize, + numReaders, + phaser, + callback, + () -> { + PaginatedEntityTimeSeriesSource source = + (timeSeriesStartTs != null) + ? new PaginatedEntityTimeSeriesSource( + entityType, + batchSize, + getSearchIndexFields(entityType), + totalRecords, + timeSeriesStartTs, + timeSeriesEndTs) + : new PaginatedEntityTimeSeriesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + return source::readWithCursor; + }, + (readers, total) -> { + List cursors = new ArrayList<>(); + int perReader = total / readers; + for (int i = 1; i < readers; i++) { + cursors.add(RestUtil.encodeCursor(String.valueOf(i * perReader))); + } + return cursors; + }); + } else { + PaginatedEntitiesSource entSource = + new PaginatedEntitiesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + submitReaders( + entityType, + totalRecords, + batchSize, + numReaders, + phaser, + callback, + () -> { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + return source::readNextKeyset; + }, + entSource::findBoundaryCursors); + } + } catch (Exception e) { + LOG.error( + "Failed to submit readers for {}, deregistering {} phaser parties", + entityType, + numReaders, + e); + for (int i = 0; i < numReaders; i++) { + phaser.arriveAndDeregister(); + } + throw e; + } + + return numReaders; + } + + public void stop() { + stopped.set(true); + } + + @Override + public void close() { + stop(); + } + + private void submitReaders( + String entityType, + int totalRecords, + int batchSize, + int numReaders, + Phaser phaser, + BatchCallback callback, + java.util.function.Supplier readerFactory, + BoundaryFinder boundaryFinder) { + if (numReaders == 1) { + KeysetBatchReader reader = readerFactory.get(); + producerExecutor.submit( + () -> + readKeysetBatches( + entityType, Integer.MAX_VALUE, batchSize, null, reader, phaser, callback)); + return; + } + + List boundaries = boundaryFinder.findBoundaries(numReaders, totalRecords); + int actualReaders = boundaries.size() + 1; + int recordsPerReader = (totalRecords + actualReaders - 1) / actualReaders; + + if (actualReaders < numReaders) { + LOG.warn( + "Boundary discovery for {} returned {} cursors (expected {}), using {} readers", + entityType, + boundaries.size(), + numReaders - 1, + actualReaders); + for (int j = 0; j < numReaders - actualReaders; j++) { + phaser.arriveAndDeregister(); + } + } + + for (int i = 0; i < actualReaders; i++) { + String startCursor = (i == 0) ? null : boundaries.get(i - 1); + int limit = (i == actualReaders - 1) ? Integer.MAX_VALUE : recordsPerReader; + KeysetBatchReader readerSource = readerFactory.get(); + final int readerLimit = limit; + producerExecutor.submit( + () -> + readKeysetBatches( + entityType, readerLimit, batchSize, startCursor, readerSource, phaser, callback)); + } + } + + private void readKeysetBatches( + String entityType, + int recordLimit, + int batchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser phaser, + BatchCallback callback) { + try { + String keysetCursor = startCursor; + int processed = 0; + + while (processed < recordLimit && !stopped.get()) { + ResultList result = readWithRetry(batchReader, keysetCursor, entityType); + if (stopped.get()) { + break; + } + + if (result == null || result.getData().isEmpty()) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (empty result)", + entityType, + processed, + recordLimit); + break; + } + + callback.onBatchRead(entityType, result, processed); + + int readCount = result.getData().size(); + int errorCount = result.getErrors() != null ? result.getErrors().size() : 0; + int warningsCount = result.getWarningsCount() != null ? result.getWarningsCount() : 0; + processed += readCount + errorCount + warningsCount; + + keysetCursor = result.getPaging() != null ? result.getPaging().getAfter() : null; + if (keysetCursor == null) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (null cursor)", + entityType, + processed, + recordLimit); + break; + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warn("Interrupted during reading of {}", entityType); + } catch (SearchIndexException e) { + LOG.error("Error reading keyset batch for {}", entityType, e); + } catch (Exception e) { + if (!stopped.get()) { + LOG.error("Error in keyset reading for {}", entityType, e); + } + } finally { + phaser.arriveAndDeregister(); + } + } + + private ResultList readWithRetry( + KeysetBatchReader batchReader, String keysetCursor, String entityType) + throws SearchIndexException, InterruptedException { + for (int attempt = 0; attempt <= maxRetryAttempts; attempt++) { + try { + return batchReader.readNextKeyset(keysetCursor); + } catch (SearchIndexException e) { + if (attempt >= maxRetryAttempts || !isTransientError(e)) { + throw e; + } + long backoff = retryBackoffMs * (1L << attempt); + LOG.warn( + "Transient read failure for {} (attempt {}/{}), retrying in {}ms", + entityType, + attempt + 1, + maxRetryAttempts, + backoff); + Thread.sleep(Math.min(backoff, 10_000)); + } + } + return null; + } + + static boolean isTransientError(SearchIndexException e) { + String msg = e.getMessage(); + if (msg == null) { + return false; + } + String lower = msg.toLowerCase(); + return lower.contains("timeout") + || lower.contains("connection") + || lower.contains("pool exhausted") + || lower.contains("connectexception") + || lower.contains("sockettimeoutexception"); + } + + static List getSearchIndexFields(String entityType) { + if (TIME_SERIES_ENTITIES.contains(entityType)) { + return List.of(); + } + return List.of("*"); + } + + static int calculateNumberOfReaders(int totalEntityRecords, int batchSize) { + if (batchSize <= 0) return 1; + return (totalEntityRecords + batchSize - 1) / batchSize; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ExecutionResult.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ExecutionResult.java index eca180a7d04..eae7af2290d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ExecutionResult.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ExecutionResult.java @@ -1,5 +1,8 @@ package org.openmetadata.service.apps.bundles.searchIndex; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; import org.openmetadata.schema.system.Stats; /** @@ -13,52 +16,61 @@ public record ExecutionResult( long failedRecords, long startTime, long endTime, - Stats finalStats) { + Stats finalStats, + Map metadata) { + + public ExecutionResult( + Status status, + long totalRecords, + long successRecords, + long failedRecords, + long startTime, + long endTime, + Stats finalStats) { + this( + status, + totalRecords, + successRecords, + failedRecords, + startTime, + endTime, + finalStats, + Collections.emptyMap()); + } /** Execution status values */ public enum Status { - /** Job completed successfully with all records processed */ COMPLETED, - /** Job completed but some records failed */ COMPLETED_WITH_ERRORS, - /** Job failed due to an exception */ FAILED, - /** Job was stopped by user request */ STOPPED } - /** Get the duration of the execution in milliseconds */ public long getDurationMillis() { return endTime - startTime; } - /** Get the duration of the execution in seconds */ public long getDurationSeconds() { return getDurationMillis() / 1000; } - /** Get the success rate as a percentage (0-100) */ public double getSuccessRate() { return totalRecords > 0 ? (successRecords * 100.0) / totalRecords : 0; } - /** Get the processing rate in records per second */ public double getRecordsPerSecond() { long durationSeconds = getDurationSeconds(); return durationSeconds > 0 ? (double) successRecords / durationSeconds : 0; } - /** Check if the execution was successful (no failures) */ public boolean isSuccessful() { return status == Status.COMPLETED; } - /** Check if the execution completed (regardless of errors) */ public boolean isCompleted() { return status == Status.COMPLETED || status == Status.COMPLETED_WITH_ERRORS; } - /** Builder for creating ExecutionResult instances */ public static Builder builder() { return new Builder(); } @@ -71,6 +83,7 @@ public record ExecutionResult( private long startTime; private long endTime; private Stats finalStats; + private Map metadata = new HashMap<>(); public Builder status(Status status) { this.status = status; @@ -107,9 +120,26 @@ public record ExecutionResult( return this; } + public Builder metadata(Map metadata) { + this.metadata = metadata != null ? metadata : new HashMap<>(); + return this; + } + + public Builder addMetadata(String key, Object value) { + this.metadata.put(key, value); + return this; + } + public ExecutionResult build() { return new ExecutionResult( - status, totalRecords, successRecords, failedRecords, startTime, endTime, finalStats); + status, + totalRecords, + successRecords, + failedRecords, + startTime, + endTime, + finalStats, + Collections.unmodifiableMap(metadata)); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingFailureRecorder.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingFailureRecorder.java index a68bf4395fe..d42f68fa413 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingFailureRecorder.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingFailureRecorder.java @@ -55,6 +55,11 @@ public class IndexingFailureRecorder implements AutoCloseable { stackTrace); } + public void recordReaderEntityFailure( + String entityType, String entityId, String entityFqn, String errorMessage) { + recordFailure(entityType, entityId, entityFqn, FailureStage.READER, errorMessage, null); + } + public void recordSinkFailure( String entityType, String entityId, String entityFqn, String errorMessage) { recordSinkFailure(entityType, entityId, entityFqn, errorMessage, null); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java new file mode 100644 index 00000000000..85eeb11f217 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java @@ -0,0 +1,534 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; + +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.EntityTimeSeriesInterface; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.workflows.searchIndex.ReindexingUtil; +import org.slf4j.MDC; + +/** + * Quartz-decoupled indexing pipeline that orchestrates: entity discovery -> reader -> queue -> sink. + * This class can be used by SearchIndexExecutor, CLI tools, REST APIs, or unit tests. + */ +@Slf4j +public class IndexingPipeline implements AutoCloseable { + + private static final String POISON_PILL = "__POISON_PILL__"; + private static final int DEFAULT_QUEUE_SIZE = 20000; + private static final int MAX_CONSUMER_THREADS = 20; + private static final int MAX_JOB_THREADS = 30; + private static final String ENTITY_TYPE_KEY = "entityType"; + private static final String RECREATE_INDEX = "recreateIndex"; + + private final SearchRepository searchRepository; + private final CompositeProgressListener listeners; + private final AtomicBoolean stopped = new AtomicBoolean(false); + @Getter private final AtomicReference stats = new AtomicReference<>(); + + private BulkSink searchIndexSink; + private RecreateIndexHandler recreateIndexHandler; + private ReindexContext recreateContext; + private EntityReader entityReader; + private ExecutorService consumerExecutor; + private ExecutorService producerExecutor; + private ExecutorService jobExecutor; + private BlockingQueue> taskQueue; + private final Set promotedEntities = java.util.concurrent.ConcurrentHashMap.newKeySet(); + + record IndexingTask(String entityType, ResultList entities, int offset) {} + + public IndexingPipeline(SearchRepository searchRepository) { + this.searchRepository = searchRepository; + this.listeners = new CompositeProgressListener(); + } + + public IndexingPipeline addListener(ReindexingProgressListener listener) { + listeners.addListener(listener); + return this; + } + + public ExecutionResult execute( + ReindexingConfiguration config, + ReindexingJobContext context, + Set entities, + BulkSink sink, + RecreateIndexHandler handler, + ReindexContext recreateCtx) { + this.searchIndexSink = sink; + this.recreateIndexHandler = handler; + this.recreateContext = recreateCtx; + long startTime = System.currentTimeMillis(); + + stats.set(initializeStats(entities)); + listeners.onJobStarted(context); + + try { + runPipeline(config, entities); + closeSink(); + finalizeReindex(); + return buildResult(startTime); + } catch (Exception e) { + LOG.error("Pipeline execution failed", e); + listeners.onJobFailed(stats.get(), e); + return ExecutionResult.fromStats(stats.get(), ExecutionResult.Status.FAILED, startTime); + } + } + + private void runPipeline(ReindexingConfiguration config, Set entities) + throws InterruptedException { + int numConsumers = + config.consumerThreads() > 0 ? Math.min(config.consumerThreads(), MAX_CONSUMER_THREADS) : 2; + int queueSize = config.queueSize() > 0 ? config.queueSize() : DEFAULT_QUEUE_SIZE; + int batchSize = config.batchSize(); + + taskQueue = new LinkedBlockingQueue<>(queueSize); + consumerExecutor = + Executors.newFixedThreadPool( + numConsumers, Thread.ofPlatform().name("pipeline-consumer-", 0).factory()); + producerExecutor = + Executors.newFixedThreadPool( + config.producerThreads() > 0 ? config.producerThreads() : 2, + Thread.ofPlatform().name("pipeline-producer-", 0).factory()); + jobExecutor = + Executors.newFixedThreadPool( + Math.min(entities.size(), MAX_JOB_THREADS), + Thread.ofPlatform().name("pipeline-job-", 0).factory()); + + entityReader = new EntityReader(producerExecutor, stopped); + + CountDownLatch consumerLatch = new CountDownLatch(numConsumers); + Map mdc = MDC.getCopyOfContextMap(); + for (int i = 0; i < numConsumers; i++) { + final int id = i; + consumerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + runConsumer(id, consumerLatch); + } finally { + MDC.clear(); + } + }); + } + + try { + readAllEntities(config, entities, batchSize); + signalConsumersToStop(numConsumers); + consumerLatch.await(); + } catch (InterruptedException e) { + stopped.set(true); + Thread.currentThread().interrupt(); + throw e; + } finally { + shutdownExecutors(); + } + } + + private void readAllEntities(ReindexingConfiguration config, Set entities, int batchSize) + throws InterruptedException { + List ordered = EntityPriority.sortByPriority(entities); + Phaser producerPhaser = new Phaser(entities.size()); + Map mdc = MDC.getCopyOfContextMap(); + + for (String entityType : ordered) { + jobExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + int totalRecords = getTotalEntityRecords(entityType); + listeners.onEntityTypeStarted(entityType, totalRecords); + + int effectiveBatchSize = + EntityBatchSizeEstimator.estimateBatchSize(entityType, batchSize); + Long filterStartTs = null; + Long filterEndTs = null; + long startTs = config.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + filterStartTs = startTs; + filterEndTs = System.currentTimeMillis(); + } + entityReader.readEntity( + entityType, + totalRecords, + effectiveBatchSize, + producerPhaser, + (type, batch, offset) -> { + if (!stopped.get()) { + taskQueue.put(new IndexingTask<>(type, batch, offset)); + } + }, + filterStartTs, + filterEndTs); + } catch (Exception e) { + LOG.error("Error reading entity type {}", entityType, e); + } finally { + producerPhaser.arriveAndDeregister(); + MDC.clear(); + } + }); + } + + int phase = 0; + while (!producerPhaser.isTerminated()) { + if (stopped.get() || Thread.currentThread().isInterrupted()) { + break; + } + try { + producerPhaser.awaitAdvanceInterruptibly(phase, 1, TimeUnit.SECONDS); + break; + } catch (TimeoutException e) { + // Continue + } + } + } + + @SuppressWarnings("unchecked") + private void runConsumer(int consumerId, CountDownLatch consumerLatch) { + try { + while (!stopped.get()) { + IndexingTask task = taskQueue.poll(200, TimeUnit.MILLISECONDS); + if (task == null) continue; + if (POISON_PILL.equals(task.entityType())) break; + + String entityType = task.entityType(); + ResultList entities = task.entities(); + Map contextData = createContextData(entityType); + + int readerSuccess = listOrEmpty(entities.getData()).size(); + int readerFailed = listOrEmpty(entities.getErrors()).size(); + int readerWarnings = entities.getWarningsCount() != null ? entities.getWarningsCount() : 0; + updateReaderStats(readerSuccess, readerFailed, readerWarnings); + + try { + if (!EntityReader.TIME_SERIES_ENTITIES.contains(entityType)) { + searchIndexSink.write( + (java.util.List) entities.getData(), contextData); + } else { + searchIndexSink.write( + (java.util.List) entities.getData(), contextData); + } + + StepStats entityStats = new StepStats(); + entityStats.setSuccessRecords(readerSuccess); + entityStats.setFailedRecords(readerFailed); + updateEntityAndJobStats(entityType, entityStats); + listeners.onProgressUpdate(stats.get(), null); + } catch (Exception e) { + LOG.error("Sink error for {}", entityType, e); + IndexingError error = + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.SINK) + .withMessage(e.getMessage()); + listeners.onError(entityType, error, stats.get()); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + consumerLatch.countDown(); + } + } + + private Map createContextData(String entityType) { + Map contextData = new HashMap<>(); + contextData.put(ENTITY_TYPE_KEY, entityType); + contextData.put(RECREATE_INDEX, recreateContext != null); + if (recreateContext != null) { + contextData.put(ReindexingUtil.RECREATE_CONTEXT, recreateContext); + recreateContext + .getStagedIndex(entityType) + .ifPresent(index -> contextData.put(ReindexingUtil.TARGET_INDEX_KEY, index)); + } + return contextData; + } + + private void signalConsumersToStop(int numConsumers) throws InterruptedException { + for (int i = 0; i < numConsumers; i++) { + taskQueue.put(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + + private void closeSink() throws IOException { + if (searchIndexSink != null) { + int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); + if (pendingVectorTasks > 0) { + LOG.info("Waiting for {} pending vector embedding tasks", pendingVectorTasks); + VectorCompletionResult vcResult = searchIndexSink.awaitVectorCompletionWithDetails(300); + LOG.info( + "Vector completion: completed={}, pending={}, waited={}ms", + vcResult.completed(), + vcResult.pendingTaskCount(), + vcResult.waitedMillis()); + } + searchIndexSink.close(); + syncSinkStats(); + } + } + + private void finalizeReindex() { + if (recreateIndexHandler == null || recreateContext == null) return; + + try { + recreateContext + .getEntities() + .forEach( + entityType -> { + if (promotedEntities.contains(entityType)) return; + try { + EntityReindexContext ctx = buildEntityReindexContext(entityType); + recreateIndexHandler.finalizeReindex(ctx, !stopped.get()); + } catch (Exception ex) { + LOG.error("Failed to finalize reindex for {}", entityType, ex); + } + }); + } finally { + recreateContext = null; + promotedEntities.clear(); + } + } + + private EntityReindexContext buildEntityReindexContext(String entityType) { + return EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) + .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases( + new HashSet<>( + org.openmetadata.common.utils.CommonUtil.listOrEmpty( + recreateContext.getParentAliases(entityType)))) + .build(); + } + + private ExecutionResult buildResult(long startTime) { + syncSinkStats(); + Stats currentStats = stats.get(); + if (currentStats != null) { + StatsReconciler.reconcile(currentStats); + } + + ExecutionResult.Status status; + if (stopped.get()) { + status = ExecutionResult.Status.STOPPED; + listeners.onJobStopped(currentStats); + } else if (hasFailures()) { + status = ExecutionResult.Status.COMPLETED_WITH_ERRORS; + listeners.onJobCompletedWithErrors(currentStats, System.currentTimeMillis() - startTime); + } else { + status = ExecutionResult.Status.COMPLETED; + listeners.onJobCompleted(currentStats, System.currentTimeMillis() - startTime); + } + + return ExecutionResult.fromStats(currentStats, status, startTime); + } + + private boolean hasFailures() { + Stats s = stats.get(); + if (s == null || s.getJobStats() == null) return false; + StepStats js = s.getJobStats(); + long failed = js.getFailedRecords() != null ? js.getFailedRecords() : 0; + long success = js.getSuccessRecords() != null ? js.getSuccessRecords() : 0; + long total = js.getTotalRecords() != null ? js.getTotalRecords() : 0; + return failed > 0 || (total > 0 && success < total); + } + + private Stats initializeStats(Set entities) { + Stats s = new Stats(); + s.setEntityStats(new org.openmetadata.schema.system.EntityStats()); + s.setJobStats(new StepStats()); + s.setReaderStats(new StepStats()); + s.setSinkStats(new StepStats()); + + int total = 0; + for (String entityType : entities) { + int entityTotal = getTotalEntityRecords(entityType); + total += entityTotal; + StepStats es = new StepStats(); + es.setTotalRecords(entityTotal); + es.setSuccessRecords(0); + es.setFailedRecords(0); + s.getEntityStats().getAdditionalProperties().put(entityType, es); + } + s.getJobStats().setTotalRecords(total); + s.getJobStats().setSuccessRecords(0); + s.getJobStats().setFailedRecords(0); + s.getReaderStats().setTotalRecords(total); + s.getReaderStats().setSuccessRecords(0); + s.getReaderStats().setFailedRecords(0); + s.getReaderStats().setWarningRecords(0); + s.getSinkStats().setTotalRecords(0); + s.getSinkStats().setSuccessRecords(0); + s.getSinkStats().setFailedRecords(0); + + s.setProcessStats(new StepStats()); + s.getProcessStats().setTotalRecords(0); + s.getProcessStats().setSuccessRecords(0); + s.getProcessStats().setFailedRecords(0); + return s; + } + + private int getTotalEntityRecords(String entityType) { + StepStats es = + stats.get() != null + && stats.get().getEntityStats() != null + && stats.get().getEntityStats().getAdditionalProperties() != null + ? stats.get().getEntityStats().getAdditionalProperties().get(entityType) + : null; + if (es != null && es.getTotalRecords() != null) { + return es.getTotalRecords(); + } + return 0; + } + + private synchronized void updateReaderStats(int success, int failed, int warnings) { + Stats s = stats.get(); + if (s == null) return; + StepStats rs = s.getReaderStats(); + if (rs == null) { + rs = new StepStats(); + s.setReaderStats(rs); + } + rs.setSuccessRecords((rs.getSuccessRecords() != null ? rs.getSuccessRecords() : 0) + success); + rs.setFailedRecords((rs.getFailedRecords() != null ? rs.getFailedRecords() : 0) + failed); + rs.setWarningRecords((rs.getWarningRecords() != null ? rs.getWarningRecords() : 0) + warnings); + } + + private synchronized void updateEntityAndJobStats(String entityType, StepStats entityDelta) { + Stats s = stats.get(); + if (s == null || s.getEntityStats() == null) return; + + StepStats es = s.getEntityStats().getAdditionalProperties().get(entityType); + if (es != null) { + es.setSuccessRecords(es.getSuccessRecords() + entityDelta.getSuccessRecords()); + es.setFailedRecords(es.getFailedRecords() + entityDelta.getFailedRecords()); + } + + StepStats js = s.getJobStats(); + if (js != null) { + int totalSuccess = + s.getEntityStats().getAdditionalProperties().values().stream() + .mapToInt(StepStats::getSuccessRecords) + .sum(); + int totalFailed = + s.getEntityStats().getAdditionalProperties().values().stream() + .mapToInt(StepStats::getFailedRecords) + .sum(); + js.setSuccessRecords(totalSuccess); + js.setFailedRecords(totalFailed); + } + } + + private synchronized void syncSinkStats() { + if (searchIndexSink == null) return; + Stats s = stats.get(); + if (s == null) return; + + StepStats bulkStats = searchIndexSink.getStats(); + if (bulkStats == null) return; + + StepStats sinkStats = s.getSinkStats(); + if (sinkStats == null) { + sinkStats = new StepStats(); + s.setSinkStats(sinkStats); + } + sinkStats.setTotalRecords( + bulkStats.getTotalRecords() != null ? bulkStats.getTotalRecords() : 0); + sinkStats.setSuccessRecords( + bulkStats.getSuccessRecords() != null ? bulkStats.getSuccessRecords() : 0); + sinkStats.setFailedRecords( + bulkStats.getFailedRecords() != null ? bulkStats.getFailedRecords() : 0); + + StepStats vectorStats = searchIndexSink.getVectorStats(); + if (vectorStats != null + && vectorStats.getTotalRecords() != null + && vectorStats.getTotalRecords() > 0) { + s.setVectorStats(vectorStats); + } + + StepStats processStats = searchIndexSink.getProcessStats(); + if (processStats != null) { + s.setProcessStats(processStats); + } + } + + private void shutdownExecutors() { + shutdownExecutor(producerExecutor, "producer"); + shutdownExecutor(jobExecutor, "job"); + shutdownExecutor(consumerExecutor, "consumer"); + } + + private void shutdownExecutor(ExecutorService executor, String name) { + if (executor != null && !executor.isShutdown()) { + executor.shutdown(); + try { + if (!executor.awaitTermination(30, TimeUnit.SECONDS)) { + executor.shutdownNow(); + LOG.warn("{} executor did not terminate in time", name); + } + } catch (InterruptedException e) { + executor.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + } + + public void stop() { + stopped.set(true); + if (entityReader != null) entityReader.stop(); + + if (searchIndexSink != null) { + LOG.info( + "Stopping pipeline: flushing sink ({} active bulk requests)", + searchIndexSink.getActiveBulkRequestCount()); + searchIndexSink.flushAndAwait(10); + } + + int dropped = taskQueue != null ? taskQueue.size() : 0; + if (dropped > 0) { + LOG.warn("Dropping {} queued tasks during shutdown", dropped); + } + + if (taskQueue != null) { + taskQueue.clear(); + for (int i = 0; i < MAX_CONSUMER_THREADS; i++) { + taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + shutdownExecutors(); + } + + @Override + public void close() { + stop(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java new file mode 100644 index 00000000000..e7d4b2018b9 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java @@ -0,0 +1,21 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Optional; +import org.openmetadata.schema.system.Stats; + +/** + * Strategy interface for reindexing execution. Encapsulates the differences between single-server + * and distributed indexing so that SearchIndexApp uses a single code path regardless of mode. + */ +public interface IndexingStrategy { + + void addListener(ReindexingProgressListener listener); + + ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context); + + Optional getStats(); + + void stop(); + + boolean isStopped(); +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java index 0c4f7025d3b..fc29a6a2aa5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java @@ -86,6 +86,10 @@ public class OpenSearchBulkSink implements BulkSink { private final AtomicLong totalSuccess = new AtomicLong(0); private final AtomicLong totalFailed = new AtomicLong(0); + // Process stage metrics (document building/transformation) + private final AtomicLong processSuccess = new AtomicLong(0); + private final AtomicLong processFailed = new AtomicLong(0); + // Configuration private volatile int batchSize; private volatile int maxConcurrentRequests; @@ -134,6 +138,7 @@ public class OpenSearchBulkSink implements BulkSink { concurrentRequests, maxPayloadSizeBytes / (1024 * 1024)); + BulkCircuitBreaker circuitBreaker = new BulkCircuitBreaker(5, 30_000, 10_000); return new CustomBulkProcessor( searchClient, bulkActions, @@ -145,7 +150,8 @@ public class OpenSearchBulkSink implements BulkSink { totalSubmitted, totalSuccess, totalFailed, - this::updateStats); + this::updateStats, + circuitBreaker); } @Override @@ -305,12 +311,14 @@ public class OpenSearchBulkSink implements BulkSink { tracker.incrementPendingSink(); } bulkProcessor.add(operation, docId, entityType, tracker, estimatedSize); + processSuccess.incrementAndGet(); if (tracker != null) { tracker.recordProcess(StatsResult.SUCCESS); } } catch (EntityNotFoundException e) { LOG.error("Entity Not Found Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -321,12 +329,14 @@ public class OpenSearchBulkSink implements BulkSink { entityTypeName, entity.getId() != null ? entity.getId().toString() : null, entity.getFullyQualifiedName(), - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } catch (Exception e) { LOG.error( "Encountered Issue while building SearchDoc from Entity Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -337,7 +347,8 @@ public class OpenSearchBulkSink implements BulkSink { entityTypeName, entity.getId() != null ? entity.getId().toString() : null, entity.getFullyQualifiedName(), - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } } @@ -364,12 +375,14 @@ public class OpenSearchBulkSink implements BulkSink { tracker.incrementPendingSink(); } bulkProcessor.add(operation, docId, entityType, tracker, estimatedSize); + processSuccess.incrementAndGet(); if (tracker != null) { tracker.recordProcess(StatsResult.SUCCESS); } } catch (EntityNotFoundException e) { LOG.error("Entity Not Found Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -379,12 +392,14 @@ public class OpenSearchBulkSink implements BulkSink { entityType, entity.getId() != null ? entity.getId().toString() : null, null, - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } catch (Exception e) { LOG.error( "Encountered Issue while building SearchDoc from Entity Due to : {}", e.getMessage(), e); totalFailed.incrementAndGet(); + processFailed.incrementAndGet(); updateStats(); if (tracker != null) { tracker.recordProcess(StatsResult.FAILED); @@ -394,7 +409,8 @@ public class OpenSearchBulkSink implements BulkSink { entityType, entity.getId() != null ? entity.getId().toString() : null, null, - e.getMessage()); + e.getMessage(), + IndexingFailureRecorder.FailureStage.PROCESS); } } } @@ -649,6 +665,28 @@ public class OpenSearchBulkSink implements BulkSink { }); } + @Override + public int getActiveBulkRequestCount() { + return bulkProcessor.activeBulkRequests.get(); + } + + @Override + public VectorCompletionResult awaitVectorCompletionWithDetails(int timeoutSeconds) { + long start = System.currentTimeMillis(); + boolean ok = awaitVectorCompletion(timeoutSeconds); + long waited = System.currentTimeMillis() - start; + if (!ok) { + int pending = getPendingVectorTaskCount(); + LOG.warn("Vector completion timed out with {} pending tasks after {}ms", pending, waited); + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordVectorTimeout(pending); + } + return VectorCompletionResult.timeout(pending, waited); + } + return VectorCompletionResult.success(waited); + } + @Override public boolean awaitVectorCompletion(int timeoutSeconds) { try { @@ -677,6 +715,16 @@ public class OpenSearchBulkSink implements BulkSink { .withFailedRecords((int) vectorFailed.get()); } + @Override + public StepStats getProcessStats() { + long success = processSuccess.get(); + long failed = processFailed.get(); + return new StepStats() + .withTotalRecords((int) (success + failed)) + .withSuccessRecords((int) success) + .withFailedRecords((int) failed); + } + public static class CustomBulkProcessor { private final OpenSearchAsyncClient asyncClient; private final List buffer = new ArrayList<>(); @@ -705,6 +753,7 @@ public class OpenSearchBulkSink implements BulkSink { private volatile boolean closed = false; private volatile FailureCallback failureCallback; private volatile SinkStatsCallback statsCallback; + private final BulkCircuitBreaker circuitBreaker; CustomBulkProcessor( OpenSearchClient client, @@ -717,7 +766,8 @@ public class OpenSearchBulkSink implements BulkSink { AtomicLong totalSubmitted, AtomicLong totalSuccess, AtomicLong totalFailed, - Runnable statsUpdater) { + Runnable statsUpdater, + BulkCircuitBreaker circuitBreaker) { this.asyncClient = new OpenSearchAsyncClient(client.getNewClient()._transport()); this.bulkActions = bulkActions; this.maxPayloadSizeBytes = maxPayloadSizeBytes; @@ -728,6 +778,7 @@ public class OpenSearchBulkSink implements BulkSink { this.totalSuccess = totalSuccess; this.totalFailed = totalFailed; this.statsUpdater = statsUpdater; + this.circuitBreaker = circuitBreaker; this.scheduler = Executors.newScheduledThreadPool(1); scheduler.scheduleAtFixedRate( @@ -852,9 +903,16 @@ public class OpenSearchBulkSink implements BulkSink { } List toFlush = new ArrayList<>(buffer); + long payloadSize = currentBufferSize; buffer.clear(); currentBufferSize = 0; + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordPayloadSize(payloadSize); + metrics.incrementPendingBulkRequests(); + } + long executionId = executionIdCounter.incrementAndGet(); int numberOfActions = toFlush.size(); totalSubmitted.addAndGet(numberOfActions); @@ -876,22 +934,69 @@ public class OpenSearchBulkSink implements BulkSink { private void executeBulkWithRetry( List operations, long executionId, int numberOfActions, int attemptNumber) { + if (!circuitBreaker.allowRequest()) { + LOG.warn( + "Circuit breaker OPEN - fail-fast for bulk request {} with {} actions", + executionId, + numberOfActions); + totalFailed.addAndGet(numberOfActions); + statsUpdater.run(); + activeBulkRequests.decrementAndGet(); + concurrentRequestSemaphore.release(); + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.decrementPendingBulkRequests(); + } + return; + } + + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + io.micrometer.core.instrument.Timer.Sample bulkTimerSample = + metrics != null ? metrics.startBulkRequestTimer() : null; + CompletableFuture future; try { future = asyncClient.bulk(b -> b.operations(operations).refresh(Refresh.False)); } catch (IOException e) { - handleBulkFailure(operations, executionId, numberOfActions, attemptNumber, e); + if (metrics != null && bulkTimerSample != null) { + metrics.recordBulkRequestCompleted(bulkTimerSample, false); + } + circuitBreaker.recordFailure(); + boolean retryScheduled = + handleBulkFailure(operations, executionId, numberOfActions, attemptNumber, e); + if (!retryScheduled) { + activeBulkRequests.decrementAndGet(); + concurrentRequestSemaphore.release(); + if (metrics != null) { + metrics.decrementPendingBulkRequests(); + } + } return; } future.whenComplete( (response, error) -> { + boolean retryScheduled = false; try { if (error != null) { - handleBulkFailure(operations, executionId, numberOfActions, attemptNumber, error); + if (metrics != null && bulkTimerSample != null) { + metrics.recordBulkRequestCompleted(bulkTimerSample, false); + } + circuitBreaker.recordFailure(); + retryScheduled = + handleBulkFailure( + operations, executionId, numberOfActions, attemptNumber, error); } else if (response.errors()) { + if (metrics != null && bulkTimerSample != null) { + metrics.recordBulkRequestCompleted(bulkTimerSample, false); + } + circuitBreaker.recordSuccess(); handlePartialFailure(response, executionId, numberOfActions); } else { + if (metrics != null && bulkTimerSample != null) { + metrics.recordBulkRequestCompleted(bulkTimerSample, true); + } + circuitBreaker.recordSuccess(); totalSuccess.addAndGet(numberOfActions); LOG.debug( "Bulk request {} completed successfully with {} actions", @@ -901,23 +1006,25 @@ public class OpenSearchBulkSink implements BulkSink { statsUpdater.run(); } } finally { - if (error != null && shouldRetry(attemptNumber, error)) { - // Don't release resources yet, we're retrying - } else { + if (!retryScheduled) { activeBulkRequests.decrementAndGet(); concurrentRequestSemaphore.release(); + if (metrics != null) { + metrics.decrementPendingBulkRequests(); + } } } }); } - private void handleBulkFailure( + private boolean handleBulkFailure( List operations, long executionId, int numberOfActions, int attemptNumber, Throwable error) { - if (shouldRetry(attemptNumber, error)) { + if (shouldRetry(attemptNumber, error) + && circuitBreaker.getState() != BulkCircuitBreaker.State.OPEN) { long backoffTime = calculateBackoff(attemptNumber); LOG.warn( "Bulk request {} failed (attempt {}), retrying in {}ms: {}", @@ -930,6 +1037,7 @@ public class OpenSearchBulkSink implements BulkSink { () -> executeBulkWithRetry(operations, executionId, numberOfActions, attemptNumber + 1), backoffTime, TimeUnit.MILLISECONDS); + return true; } else { totalFailed.addAndGet(numberOfActions); LOG.error( @@ -955,7 +1063,12 @@ public class OpenSearchBulkSink implements BulkSink { tracker.recordSink(StatsResult.FAILED); } if (failureCallback != null) { - failureCallback.onFailure(entityType, docId, null, error.getMessage()); + failureCallback.onFailure( + entityType, + docId, + null, + error.getMessage(), + IndexingFailureRecorder.FailureStage.SINK); } } } @@ -965,6 +1078,7 @@ public class OpenSearchBulkSink implements BulkSink { } } statsUpdater.run(); + return false; } } @@ -996,7 +1110,8 @@ public class OpenSearchBulkSink implements BulkSink { tracker.recordSink(StatsResult.FAILED); } if (failureCallback != null) { - failureCallback.onFailure(entityType, docId, null, failureMessage); + failureCallback.onFailure( + entityType, docId, null, failureMessage, IndexingFailureRecorder.FailureStage.SINK); } } else { String entityType = docId != null ? docIdToEntityType.remove(docId) : null; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java new file mode 100644 index 00000000000..1fb84174d3a --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java @@ -0,0 +1,32 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Map; +import java.util.UUID; +import org.openmetadata.schema.entity.app.AppRunRecord; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.Stats; + +public interface OrchestratorContext { + + String getJobName(); + + String getAppConfigJson(); + + void storeRunStats(Stats stats); + + void storeRunRecord(String json); + + AppRunRecord getJobRecord(); + + void pushStatusUpdate(AppRunRecord record, boolean force); + + UUID getAppId(); + + Map getAppConfiguration(); + + void updateAppConfiguration(Map config); + + ReindexingProgressListener createProgressListener(EventPublisherJob jobData); + + ReindexingJobContext createReindexingContext(boolean distributed); +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java new file mode 100644 index 00000000000..497616eac5e --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java @@ -0,0 +1,98 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.service.apps.scheduler.OmAppJobListener.APP_CONFIG; +import static org.openmetadata.service.apps.scheduler.OmAppJobListener.APP_RUN_STATS; + +import java.util.Map; +import java.util.UUID; +import java.util.function.Function; +import org.openmetadata.schema.entity.app.App; +import org.openmetadata.schema.entity.app.AppRunRecord; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.apps.bundles.searchIndex.listeners.QuartzProgressListener; +import org.quartz.JobExecutionContext; + +public class QuartzOrchestratorContext implements OrchestratorContext { + + private static final String APP_SCHEDULE_RUN = "AppScheduleRun"; + + private final JobExecutionContext ctx; + private final App app; + private final Function jobRecordProvider; + private final StatusPusher statusPusher; + + @FunctionalInterface + public interface StatusPusher { + void push(JobExecutionContext ctx, AppRunRecord record, boolean force); + } + + public QuartzOrchestratorContext( + JobExecutionContext ctx, + App app, + Function jobRecordProvider, + StatusPusher statusPusher) { + this.ctx = ctx; + this.app = app; + this.jobRecordProvider = jobRecordProvider; + this.statusPusher = statusPusher; + } + + @Override + public String getJobName() { + return ctx.getJobDetail().getKey().getName(); + } + + @Override + public String getAppConfigJson() { + return (String) ctx.getJobDetail().getJobDataMap().get(APP_CONFIG); + } + + @Override + public void storeRunStats(Stats stats) { + ctx.getJobDetail().getJobDataMap().put(APP_RUN_STATS, stats); + } + + @Override + public void storeRunRecord(String json) { + ctx.getJobDetail().getJobDataMap().put(APP_SCHEDULE_RUN, json); + } + + @Override + public AppRunRecord getJobRecord() { + return jobRecordProvider.apply(ctx); + } + + @Override + public void pushStatusUpdate(AppRunRecord record, boolean force) { + statusPusher.push(ctx, record, force); + } + + @Override + public UUID getAppId() { + return app != null ? app.getId() : null; + } + + @Override + public Map getAppConfiguration() { + return app != null ? JsonUtils.getMap(app.getAppConfiguration()) : null; + } + + @Override + public void updateAppConfiguration(Map config) { + if (app != null) { + app.setAppConfiguration(config); + } + } + + @Override + public ReindexingProgressListener createProgressListener(EventPublisherJob jobData) { + return new QuartzProgressListener(ctx, jobData, app, jobRecordProvider, statusPusher); + } + + @Override + public ReindexingJobContext createReindexingContext(boolean distributed) { + return new QuartzJobContext(ctx, app, distributed); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java index e3ec6037ad8..29b5d1bf3f0 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java @@ -1,8 +1,14 @@ package org.openmetadata.service.apps.bundles.searchIndex; +import java.util.Collections; +import java.util.Map; import java.util.Set; import org.openmetadata.schema.system.EventPublisherJob; import org.openmetadata.schema.type.IndexMappingLanguage; +import org.openmetadata.service.search.SearchClusterMetrics; +import org.openmetadata.service.search.SearchRepository; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Immutable configuration for a reindexing job. This record encapsulates all the configuration @@ -26,7 +32,11 @@ public record ReindexingConfiguration( IndexMappingLanguage searchIndexMappingLanguage, String afterCursor, String slackBotToken, - String slackChannel) { + String slackChannel, + int timeSeriesMaxDays, + Map timeSeriesEntityDays) { + + private static final Logger LOG = LoggerFactory.getLogger(ReindexingConfiguration.class); private static final int DEFAULT_BATCH_SIZE = 100; private static final int DEFAULT_CONSUMER_THREADS = 1; @@ -37,6 +47,50 @@ public record ReindexingConfiguration( private static final int DEFAULT_MAX_RETRIES = 5; private static final int DEFAULT_INITIAL_BACKOFF = 1000; private static final int DEFAULT_MAX_BACKOFF = 10000; + private static final int DEFAULT_TIME_SERIES_MAX_DAYS = 0; + + public static ReindexingConfiguration applyAutoTuning( + ReindexingConfiguration config, SearchRepository searchRepository) { + if (!config.autoTune()) { + return config; + } + SearchClusterMetrics metrics = fetchClusterMetrics(searchRepository); + if (metrics == null) { + return config; + } + return ReindexingConfiguration.builder() + .entities(config.entities()) + .batchSize(metrics.getRecommendedBatchSize()) + .consumerThreads(metrics.getRecommendedConsumerThreads()) + .producerThreads(metrics.getRecommendedProducerThreads()) + .queueSize(metrics.getRecommendedQueueSize()) + .maxConcurrentRequests(metrics.getRecommendedConcurrentRequests()) + .payloadSize(metrics.getMaxPayloadSizeBytes()) + .recreateIndex(config.recreateIndex()) + .autoTune(true) + .useDistributedIndexing(config.useDistributedIndexing()) + .force(config.force()) + .maxRetries(config.maxRetries()) + .initialBackoff(config.initialBackoff()) + .maxBackoff(config.maxBackoff()) + .searchIndexMappingLanguage(config.searchIndexMappingLanguage()) + .afterCursor(config.afterCursor()) + .slackBotToken(config.slackBotToken()) + .slackChannel(config.slackChannel()) + .timeSeriesMaxDays(config.timeSeriesMaxDays()) + .timeSeriesEntityDays(config.timeSeriesEntityDays()) + .build(); + } + + private static SearchClusterMetrics fetchClusterMetrics(SearchRepository searchRepository) { + try { + return SearchClusterMetrics.fetchClusterMetrics( + searchRepository, 0, searchRepository.getMaxDBConnections()); + } catch (Exception e) { + LOG.warn("Failed to fetch cluster metrics for auto-tuning, using configured values", e); + return null; + } + } /** * Creates a ReindexingConfiguration from an EventPublisherJob. @@ -69,7 +123,30 @@ public record ReindexingConfiguration( jobData.getSearchIndexMappingLanguage(), jobData.getAfterCursor(), jobData.getSlackBotToken(), - jobData.getSlackChannel()); + jobData.getSlackChannel(), + jobData.getTimeSeriesMaxDays() != null + ? jobData.getTimeSeriesMaxDays() + : DEFAULT_TIME_SERIES_MAX_DAYS, + jobData.getTimeSeriesEntityDays() != null + ? jobData.getTimeSeriesEntityDays() + : Collections.emptyMap()); + } + + /** + * Returns the start timestamp for time series date filtering for the given entity type. Uses + * per-entity override if configured, otherwise falls back to the default timeSeriesMaxDays. + * + * @return start timestamp in millis, or -1 if no filtering should be applied (days <= 0) + */ + public long getTimeSeriesStartTs(String entityType) { + int days = timeSeriesMaxDays; + if (timeSeriesEntityDays != null && timeSeriesEntityDays.containsKey(entityType)) { + days = timeSeriesEntityDays.get(entityType); + } + if (days <= 0) { + return -1; + } + return System.currentTimeMillis() - (days * 86_400_000L); } /** Check if Slack notifications are configured */ @@ -109,6 +186,8 @@ public record ReindexingConfiguration( private String afterCursor; private String slackBotToken; private String slackChannel; + private int timeSeriesMaxDays = DEFAULT_TIME_SERIES_MAX_DAYS; + private Map timeSeriesEntityDays = Collections.emptyMap(); public Builder entities(Set entities) { this.entities = entities; @@ -200,6 +279,16 @@ public record ReindexingConfiguration( return this; } + public Builder timeSeriesMaxDays(int timeSeriesMaxDays) { + this.timeSeriesMaxDays = timeSeriesMaxDays; + return this; + } + + public Builder timeSeriesEntityDays(Map timeSeriesEntityDays) { + this.timeSeriesEntityDays = timeSeriesEntityDays; + return this; + } + public ReindexingConfiguration build() { return new ReindexingConfiguration( entities, @@ -219,7 +308,9 @@ public record ReindexingConfiguration( searchIndexMappingLanguage, afterCursor, slackBotToken, - slackChannel); + slackChannel, + timeSeriesMaxDays, + timeSeriesEntityDays); } } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetrics.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetrics.java new file mode 100644 index 00000000000..97c20abe47b --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetrics.java @@ -0,0 +1,312 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.Timer; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class ReindexingMetrics { + + private static volatile ReindexingMetrics instance; + private final MeterRegistry meterRegistry; + + // Job lifecycle + private final Counter jobsStarted; + private final Counter jobsCompleted; + private final Counter jobsFailed; + private final Counter jobsStopped; + private final Timer jobDurationCompleted; + private final Timer jobDurationFailed; + private final Timer jobDurationStopped; + private final AtomicLong activeJobs = new AtomicLong(); + + // Bulk request metrics + private final Timer bulkDurationSuccess; + private final Timer bulkDurationFailure; + private final DistributionSummary bulkPayloadSize; + private final AtomicLong pendingBulkRequests = new AtomicLong(); + + // Backpressure + private final Counter backpressureEvents; + + // Circuit breaker + private final Map circuitBreakerCounters = new ConcurrentHashMap<>(); + + // Vector timeouts + private final Counter vectorTimeouts; + + // Queue fill ratio gauge + private final AtomicLong queueFillRatio = new AtomicLong(); + + private ReindexingMetrics(MeterRegistry meterRegistry) { + this.meterRegistry = meterRegistry; + + // Job lifecycle counters + this.jobsStarted = + Counter.builder("reindexing.jobs") + .description("Job lifecycle events") + .tag("status", "started") + .register(meterRegistry); + + this.jobsCompleted = + Counter.builder("reindexing.jobs") + .description("Job lifecycle events") + .tag("status", "completed") + .register(meterRegistry); + + this.jobsFailed = + Counter.builder("reindexing.jobs") + .description("Job lifecycle events") + .tag("status", "failed") + .register(meterRegistry); + + this.jobsStopped = + Counter.builder("reindexing.jobs") + .description("Job lifecycle events") + .tag("status", "stopped") + .register(meterRegistry); + + // Job duration timers + this.jobDurationCompleted = + Timer.builder("reindexing.job.duration") + .description("Job wall-clock duration") + .tag("status", "completed") + .register(meterRegistry); + + this.jobDurationFailed = + Timer.builder("reindexing.job.duration") + .description("Job wall-clock duration") + .tag("status", "failed") + .register(meterRegistry); + + this.jobDurationStopped = + Timer.builder("reindexing.job.duration") + .description("Job wall-clock duration") + .tag("status", "stopped") + .register(meterRegistry); + + // Active jobs gauge + Gauge.builder("reindexing.jobs.active", activeJobs, AtomicLong::get) + .description("Currently running reindexing jobs") + .register(meterRegistry); + + // Bulk request timers with SLA buckets + this.bulkDurationSuccess = + Timer.builder("reindexing.bulk.duration") + .description("Bulk request latency") + .tag("success", "true") + .sla( + Duration.ofMillis(50), + Duration.ofMillis(100), + Duration.ofMillis(500), + Duration.ofSeconds(1), + Duration.ofSeconds(5), + Duration.ofSeconds(10), + Duration.ofSeconds(30)) + .register(meterRegistry); + + this.bulkDurationFailure = + Timer.builder("reindexing.bulk.duration") + .description("Bulk request latency") + .tag("success", "false") + .sla( + Duration.ofMillis(50), + Duration.ofMillis(100), + Duration.ofMillis(500), + Duration.ofSeconds(1), + Duration.ofSeconds(5), + Duration.ofSeconds(10), + Duration.ofSeconds(30)) + .register(meterRegistry); + + // Bulk payload size distribution + this.bulkPayloadSize = + DistributionSummary.builder("reindexing.bulk.payload.size") + .description("Payload size in bytes") + .baseUnit("bytes") + .serviceLevelObjectives( + 64 * 1024d, 256 * 1024d, 1024 * 1024d, 5 * 1024 * 1024d, 20 * 1024 * 1024d) + .register(meterRegistry); + + // Pending bulk requests gauge + Gauge.builder("reindexing.sink.pending", pendingBulkRequests, AtomicLong::get) + .description("In-flight bulk requests") + .register(meterRegistry); + + // Backpressure counter + this.backpressureEvents = + Counter.builder("reindexing.backpressure.events") + .description("Backpressure detections") + .register(meterRegistry); + + // Vector timeouts counter + this.vectorTimeouts = + Counter.builder("reindexing.vector.timeouts") + .description("Vector embedding completion timeouts") + .register(meterRegistry); + + // Queue fill ratio gauge + Gauge.builder("reindexing.queue.fill_ratio", queueFillRatio, AtomicLong::get) + .description("Task queue fill ratio (0-100)") + .register(meterRegistry); + + LOG.info("Reindexing metrics initialized"); + } + + public static synchronized void initialize(MeterRegistry meterRegistry) { + if (instance == null) { + instance = new ReindexingMetrics(meterRegistry); + } + } + + public static ReindexingMetrics getInstance() { + return instance; + } + + // --- Job lifecycle --- + + public Timer.Sample startJobTimer() { + return Timer.start(meterRegistry); + } + + public void recordJobStarted() { + jobsStarted.increment(); + activeJobs.incrementAndGet(); + } + + public void recordJobCompleted(Timer.Sample sample) { + jobsCompleted.increment(); + activeJobs.decrementAndGet(); + if (sample != null) { + sample.stop(jobDurationCompleted); + } + } + + public void recordJobFailed(Timer.Sample sample) { + jobsFailed.increment(); + activeJobs.decrementAndGet(); + if (sample != null) { + sample.stop(jobDurationFailed); + } + } + + public void recordJobStopped(Timer.Sample sample) { + jobsStopped.increment(); + activeJobs.decrementAndGet(); + if (sample != null) { + sample.stop(jobDurationStopped); + } + } + + // --- Stage counters (dynamic tags) --- + + public void recordStageSuccess(String stage, String entityType, long count) { + Counter.builder("reindexing.stage.success") + .description("Records successfully processed per stage") + .tag("stage", stage) + .tag("entity_type", entityType) + .register(meterRegistry) + .increment(count); + } + + public void recordStageFailed(String stage, String entityType, long count) { + Counter.builder("reindexing.stage.failed") + .description("Records failed per stage") + .tag("stage", stage) + .tag("entity_type", entityType) + .register(meterRegistry) + .increment(count); + } + + public void recordStageWarnings(String stage, String entityType, long count) { + Counter.builder("reindexing.stage.warnings") + .description("Reader warnings") + .tag("stage", stage) + .tag("entity_type", entityType) + .register(meterRegistry) + .increment(count); + } + + // --- Bulk request metrics --- + + public Timer.Sample startBulkRequestTimer() { + return Timer.start(meterRegistry); + } + + public void recordBulkRequestCompleted(Timer.Sample sample, boolean success) { + if (sample != null) { + sample.stop(success ? bulkDurationSuccess : bulkDurationFailure); + } + } + + public void recordPayloadSize(long sizeBytes) { + bulkPayloadSize.record(sizeBytes); + } + + public void incrementPendingBulkRequests() { + pendingBulkRequests.incrementAndGet(); + } + + public void decrementPendingBulkRequests() { + pendingBulkRequests.decrementAndGet(); + } + + // --- Backpressure --- + + public void recordBackpressureEvent() { + backpressureEvents.increment(); + } + + // --- Promotion metrics (dynamic tags) --- + + public void recordPromotionSuccess(String entityType) { + Counter.builder("reindexing.promotion") + .description("Index promotion events") + .tag("entity_type", entityType) + .tag("result", "success") + .register(meterRegistry) + .increment(); + } + + public void recordPromotionFailure(String entityType) { + Counter.builder("reindexing.promotion") + .description("Index promotion events") + .tag("entity_type", entityType) + .tag("result", "failure") + .register(meterRegistry) + .increment(); + } + + // --- Circuit breaker --- + + public void recordCircuitBreakerTrip(String transition) { + circuitBreakerCounters + .computeIfAbsent( + transition, + t -> + Counter.builder("reindexing.circuitbreaker.trips") + .description("Circuit breaker state transitions") + .tag("transition", t) + .register(meterRegistry)) + .increment(); + } + + // --- Vector timeouts --- + + public void recordVectorTimeout(int pendingCount) { + vectorTimeouts.increment(); + } + + // --- Queue fill ratio --- + + public void updateQueueFillRatio(int percent) { + queueFillRatio.set(percent); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java new file mode 100644 index 00000000000..44e76fe254c --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java @@ -0,0 +1,459 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.service.apps.scheduler.AppScheduler.ON_DEMAND_JOB; +import static org.openmetadata.service.socket.WebSocketManager.SEARCH_INDEX_JOB_BROADCAST_CHANNEL; + +import com.fasterxml.jackson.core.type.TypeReference; +import io.micrometer.core.instrument.Timer; +import java.util.Collections; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.api.configuration.OpenMetadataBaseUrlConfiguration; +import org.openmetadata.schema.entity.app.AppRunRecord; +import org.openmetadata.schema.entity.app.FailureContext; +import org.openmetadata.schema.entity.app.SuccessContext; +import org.openmetadata.schema.settings.Settings; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.listeners.LoggingProgressListener; +import org.openmetadata.service.apps.bundles.searchIndex.listeners.SlackProgressListener; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.SystemRepository; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.socket.WebSocketManager; +import org.slf4j.MDC; + +@Slf4j +public class ReindexingOrchestrator { + private static final String ALL = "all"; + private final CollectionDAO collectionDAO; + private final SearchRepository searchRepository; + private final OrchestratorContext context; + + @Getter private EventPublisherJob jobData; + private volatile boolean stopped = false; + private volatile IndexingStrategy activeStrategy; + private volatile Map resultMetadata = Collections.emptyMap(); + + public ReindexingOrchestrator( + CollectionDAO collectionDAO, SearchRepository searchRepository, OrchestratorContext context) { + this.collectionDAO = collectionDAO; + this.searchRepository = searchRepository; + this.context = context; + } + + public void run(EventPublisherJob initialJobData) { + this.jobData = initialJobData; + initializeState(); + initializeJobData(); + + String jobId = UUID.randomUUID().toString().substring(0, 8); + MDC.put("reindexJobId", jobId); + + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + Timer.Sample timerSample = null; + if (metrics != null) { + metrics.recordJobStarted(); + timerSample = metrics.startJobTimer(); + } + + preflightFixes(); + + try { + runReindexing(); + } catch (Exception ex) { + handleExecutionException(ex); + } finally { + finalizeJobExecution(); + cleanupOrphanedIndices(); + + if (metrics != null && timerSample != null) { + EventPublisherJob.Status status = jobData != null ? jobData.getStatus() : null; + if (status == EventPublisherJob.Status.COMPLETED + || status == EventPublisherJob.Status.ACTIVE_ERROR) { + metrics.recordJobCompleted(timerSample); + } else if (status == EventPublisherJob.Status.STOPPED) { + metrics.recordJobStopped(timerSample); + } else { + metrics.recordJobFailed(timerSample); + } + } + + MDC.remove("reindexJobId"); + } + } + + public void stop() { + LOG.info("Reindexing job is being stopped."); + stopped = true; + + IndexingStrategy strategy = this.activeStrategy; + if (strategy != null) { + try { + strategy.stop(); + } catch (Exception e) { + LOG.error("Error stopping indexing strategy", e); + } + } + + if (jobData != null) { + jobData.setStatus(EventPublisherJob.Status.STOPPED); + } + + AppRunRecord appRecord = context.getJobRecord(); + appRecord.setStatus(AppRunRecord.Status.STOPPED); + appRecord.setEndTime(System.currentTimeMillis()); + context.storeRunRecord(JsonUtils.pojoToJson(appRecord)); + context.pushStatusUpdate(appRecord, true); + sendUpdates(); + + LOG.info("Reindexing job stopped successfully."); + } + + private void initializeState() { + stopped = false; + activeStrategy = null; + resultMetadata = Collections.emptyMap(); + } + + private void initializeJobData() { + if (jobData == null) { + jobData = loadJobData(); + } + + String jobName = context.getJobName(); + if (jobName.equals(ON_DEMAND_JOB)) { + Map jsonAppConfig = + JsonUtils.convertValue(jobData, new TypeReference>() {}); + context.updateAppConfiguration(jsonAppConfig); + } + } + + private EventPublisherJob loadJobData() { + String appConfigJson = context.getAppConfigJson(); + if (appConfigJson != null) { + return JsonUtils.readValue(appConfigJson, EventPublisherJob.class); + } + + Map appConfig = context.getAppConfiguration(); + if (appConfig != null) { + return JsonUtils.convertValue(appConfig, EventPublisherJob.class); + } + + LOG.error("Unable to initialize jobData from JobDataMap or App configuration"); + throw new SearchIndexApp.ReindexingException("JobData is not initialized"); + } + + private void preflightFixes() { + LOG.info("Running preflight fixes before reindexing"); + markStaleRunningJobsStopped(); + cleanupOrphanedIndicesPreFlight(); + } + + private static final String APP_NAME = "SearchIndexingApplication"; + + private void markStaleRunningJobsStopped() { + try { + AppRunRecord currentRecord = context.getJobRecord(); + if (currentRecord != null && currentRecord.getStartTime() != null) { + collectionDAO + .appExtensionTimeSeriesDao() + .markStaleEntriesStoppedBefore(APP_NAME, currentRecord.getStartTime()); + LOG.info("Preflight: marked stale running jobs as stopped for {}", APP_NAME); + } + } catch (Exception e) { + LOG.warn("Preflight: failed to cleanup stale running jobs: {}", e.getMessage()); + } + } + + private void cleanupOrphanedIndicesPreFlight() { + try { + OrphanedIndexCleaner cleaner = new OrphanedIndexCleaner(); + OrphanedIndexCleaner.CleanupResult result = + cleaner.cleanupOrphanedIndices(searchRepository.getSearchClient()); + if (result.found() > 0) { + LOG.info( + "Preflight: cleaned up {} orphaned rebuild indices (found={}, failed={})", + result.deleted(), + result.found(), + result.failed()); + } + } catch (Exception e) { + LOG.warn("Preflight: failed to cleanup orphaned indices: {}", e.getMessage()); + } + } + + private void runReindexing() throws Exception { + if (jobData.getEntities() == null || jobData.getEntities().isEmpty()) { + LOG.info("No entities selected for reindexing, completing immediately"); + jobData.setStatus(EventPublisherJob.Status.COMPLETED); + jobData.setStats(new Stats()); + return; + } + + setupEntities(); + cleanupOldFailures(); + + LOG.info( + "Search Index Job Started for Entities: {}, RecreateIndex: {}, DistributedIndexing: {}", + jobData.getEntities(), + jobData.getRecreateIndex(), + jobData.getUseDistributedIndexing()); + + activeStrategy = createStrategy(); + + activeStrategy.addListener(context.createProgressListener(jobData)); + activeStrategy.addListener(new LoggingProgressListener()); + + if (hasSlackConfig()) { + String instanceUrl = getInstanceUrl(); + activeStrategy.addListener( + new SlackProgressListener( + jobData.getSlackBotToken(), jobData.getSlackChannel(), instanceUrl)); + } + + ReindexingJobContext jobContext = + context.createReindexingContext(Boolean.TRUE.equals(jobData.getUseDistributedIndexing())); + + ReindexingConfiguration config = ReindexingConfiguration.from(jobData); + config = ReindexingConfiguration.applyAutoTuning(config, searchRepository); + + ExecutionResult result = activeStrategy.execute(config, jobContext); + updateJobDataFromResult(result); + + if (jobData.getStats() != null) { + context.storeRunStats(jobData.getStats()); + } + + if (!result.metadata().isEmpty()) { + saveResultMetadataToJobRecord(result.metadata()); + } + } + + private IndexingStrategy createStrategy() { + if (Boolean.TRUE.equals(jobData.getUseDistributedIndexing())) { + AppRunRecord appRecord = context.getJobRecord(); + return new DistributedIndexingStrategy( + collectionDAO, + searchRepository, + jobData, + appRecord.getAppId(), + appRecord.getStartTime(), + context.getJobName()); + } + return new SingleServerIndexingStrategy(collectionDAO, searchRepository); + } + + private void updateJobDataFromResult(ExecutionResult result) { + if (result.finalStats() != null) { + Stats stats = result.finalStats(); + StatsReconciler.reconcile(stats); + jobData.setStats(stats); + } + + resultMetadata = result.metadata() != null ? result.metadata() : Collections.emptyMap(); + + switch (result.status()) { + case COMPLETED -> jobData.setStatus(EventPublisherJob.Status.COMPLETED); + case COMPLETED_WITH_ERRORS -> jobData.setStatus(EventPublisherJob.Status.ACTIVE_ERROR); + case FAILED -> jobData.setStatus(EventPublisherJob.Status.FAILED); + case STOPPED -> jobData.setStatus(EventPublisherJob.Status.STOPPED); + } + } + + private void saveResultMetadataToJobRecord(Map metadata) { + try { + AppRunRecord appRecord = context.getJobRecord(); + SuccessContext successContext = appRecord.getSuccessContext(); + if (successContext == null) { + successContext = new SuccessContext(); + } + + for (Map.Entry entry : metadata.entrySet()) { + successContext.withAdditionalProperty(entry.getKey(), entry.getValue()); + } + + if (jobData.getStats() != null) { + successContext.withAdditionalProperty("stats", jobData.getStats()); + } + + appRecord.setSuccessContext(successContext); + context.storeRunRecord(JsonUtils.pojoToJson(appRecord)); + } catch (Exception e) { + LOG.error("Failed to save result metadata to job record", e); + } + } + + private void handleExecutionException(Exception ex) { + IndexingStrategy strategy = this.activeStrategy; + if (strategy != null && jobData != null) { + try { + strategy.getStats().ifPresent(jobData::setStats); + } catch (Exception e) { + LOG.debug("Could not capture strategy stats during exception handling", e); + } + } + + if (stopped) { + if (jobData != null) { + jobData.setStatus(EventPublisherJob.Status.STOPPED); + } + } else { + IndexingError error = + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.JOB) + .withMessage("Reindexing Job Exception: " + ex.getMessage()); + LOG.error("Reindexing Job Failed", ex); + + if (jobData != null) { + jobData.setStatus(EventPublisherJob.Status.FAILED); + jobData.setFailure(error); + } + } + } + + private void finalizeJobExecution() { + sendUpdates(); + + if (stopped) { + AppRunRecord appRecord = context.getJobRecord(); + appRecord.setStatus(AppRunRecord.Status.STOPPED); + context.storeRunRecord(JsonUtils.pojoToJson(appRecord)); + } + } + + private void sendUpdates() { + try { + updateRecordToDbAndNotify(); + } catch (Exception ex) { + LOG.error("Failed to send updates", ex); + } + } + + private void updateRecordToDbAndNotify() { + AppRunRecord appRecord = context.getJobRecord(); + appRecord.setStatus(AppRunRecord.Status.fromValue(jobData.getStatus().value())); + + if (jobData.getFailure() != null) { + appRecord.setFailureContext( + new FailureContext().withAdditionalProperty("failure", jobData.getFailure())); + } + + if (jobData.getStats() != null) { + SuccessContext successContext = + new SuccessContext().withAdditionalProperty("stats", jobData.getStats()); + + String distributedJobId = (String) resultMetadata.get("distributedJobId"); + + try { + UUID appId = context.getAppId(); + String jobIdStr = + distributedJobId != null ? distributedJobId : (appId != null ? appId.toString() : null); + if (jobIdStr != null) { + int failureCount = collectionDAO.searchIndexFailureDAO().countByJobId(jobIdStr); + if (failureCount > 0) { + successContext.withAdditionalProperty("failureRecordCount", failureCount); + } + } + } catch (Exception e) { + LOG.debug("Could not get failure count", e); + } + + Object serverStats = resultMetadata.get("serverStats"); + if (serverStats != null) { + successContext.withAdditionalProperty("serverStats", serverStats); + successContext.withAdditionalProperty("serverCount", resultMetadata.get("serverCount")); + successContext.withAdditionalProperty("distributedJobId", distributedJobId); + } + + appRecord.setSuccessContext(successContext); + } + + if (WebSocketManager.getInstance() != null) { + String messageJson = JsonUtils.pojoToJson(appRecord); + WebSocketManager.getInstance() + .broadCastMessageToAll(SEARCH_INDEX_JOB_BROADCAST_CHANNEL, messageJson); + } + } + + private void cleanupOldFailures() { + try { + int deleted = collectionDAO.searchIndexFailureDAO().deleteAll(); + if (deleted > 0) { + LOG.info("Cleaned up {} failure records from previous runs", deleted); + } + } catch (Exception e) { + LOG.warn("Failed to cleanup old failure records", e); + } + } + + private void cleanupOrphanedIndices() { + try { + OrphanedIndexCleaner cleaner = new OrphanedIndexCleaner(); + OrphanedIndexCleaner.CleanupResult result = + cleaner.cleanupOrphanedIndices(searchRepository.getSearchClient()); + if (result.deleted() > 0) { + LOG.info( + "Cleaned up {} orphaned rebuild indices on Job End (found={}, failed={})", + result.deleted(), + result.found(), + result.failed()); + } + } catch (Exception e) { + LOG.warn("Failed to cleanup orphaned indices on Job End: {}", e.getMessage()); + } + } + + private void setupEntities() { + boolean containsAll = jobData.getEntities().contains(ALL); + if (containsAll) { + jobData.setEntities(getAll()); + } + } + + private Set getAll() { + Set entities = + new HashSet<>( + Entity.getEntityList().stream() + .filter(t -> searchRepository.getEntityIndexMap().containsKey(t)) + .toList()); + entities.addAll( + SearchIndexApp.TIME_SERIES_ENTITIES.stream() + .filter(t -> searchRepository.getEntityIndexMap().containsKey(t)) + .toList()); + return entities; + } + + private boolean hasSlackConfig() { + return jobData.getSlackBotToken() != null + && !jobData.getSlackBotToken().isEmpty() + && jobData.getSlackChannel() != null + && !jobData.getSlackChannel().isEmpty(); + } + + private String getInstanceUrl() { + try { + SystemRepository systemRepository = Entity.getSystemRepository(); + if (systemRepository != null) { + Settings settings = systemRepository.getOMBaseUrlConfigInternal(); + if (settings != null && settings.getConfigValue() != null) { + OpenMetadataBaseUrlConfiguration urlConfig = + (OpenMetadataBaseUrlConfiguration) settings.getConfigValue(); + if (urlConfig != null && urlConfig.getOpenMetadataUrl() != null) { + return urlConfig.getOpenMetadataUrl(); + } + } + } + } catch (Exception e) { + LOG.debug("Could not get instance URL from SystemSettings", e); + } + return "http://localhost:8585"; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java index 653b5f09e37..d04b613a3ac 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java @@ -18,6 +18,15 @@ import org.openmetadata.schema.system.StepStats; */ public interface ReindexingProgressListener { + /** Failure type classification for per-stage failure hooks */ + enum FailureType { + ENTITY_NOT_FOUND, + JSON_PROCESSING, + DB_ERROR, + SINK_ERROR, + UNKNOWN + } + /** Called when reindexing job starts initialization */ default void onJobStarted(ReindexingJobContext context) {} @@ -39,6 +48,20 @@ public interface ReindexingProgressListener { /** Called when an error occurs during processing */ default void onError(String entityType, IndexingError error, Stats currentStats) {} + /** Called when a reader-stage failure occurs for a specific entity */ + default void onReaderFailure( + String entityType, String entityId, String error, FailureType type) {} + + /** Called when a process-stage failure occurs (entity -> search doc conversion) */ + default void onProcessFailure(String entityType, String entityId, String error) {} + + /** Called when a sink-stage failure occurs (ES/OS bulk indexing) */ + default void onSinkFailure(String entityType, String entityId, String error) {} + + /** Called when sub-indexing (columns, vectors) completes for an entity type */ + default void onSubIndexingCompleted( + String entityType, String subIndex, StepStats subIndexStats) {} + /** Called when job completes successfully */ default void onJobCompleted(Stats finalStats, long elapsedMillis) {} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java index 282e4a7a2d9..3fad970caff 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java @@ -1,63 +1,24 @@ package org.openmetadata.service.apps.bundles.searchIndex; -import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; import static org.openmetadata.service.Entity.QUERY_COST_RECORD; import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; import static org.openmetadata.service.Entity.TEST_CASE_RESULT; -import static org.openmetadata.service.apps.scheduler.AppScheduler.ON_DEMAND_JOB; -import static org.openmetadata.service.apps.scheduler.OmAppJobListener.APP_CONFIG; -import static org.openmetadata.service.apps.scheduler.OmAppJobListener.APP_RUN_STATS; -import static org.openmetadata.service.socket.WebSocketManager.SEARCH_INDEX_JOB_BROADCAST_CHANNEL; -import com.fasterxml.jackson.core.type.TypeReference; import jakarta.ws.rs.core.Response; -import java.util.Collections; -import java.util.HashSet; import java.util.Map; import java.util.Set; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; import lombok.Getter; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.analytics.ReportData; -import org.openmetadata.schema.api.configuration.OpenMetadataBaseUrlConfiguration; import org.openmetadata.schema.entity.app.App; -import org.openmetadata.schema.entity.app.AppRunRecord; -import org.openmetadata.schema.entity.app.FailureContext; -import org.openmetadata.schema.entity.app.SuccessContext; -import org.openmetadata.schema.settings.Settings; import org.openmetadata.schema.system.EventPublisherJob; -import org.openmetadata.schema.system.IndexingError; -import org.openmetadata.schema.system.Stats; -import org.openmetadata.schema.system.StepStats; import org.openmetadata.schema.utils.JsonUtils; -import org.openmetadata.service.Entity; import org.openmetadata.service.apps.AbstractNativeApplication; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.DistributedSearchIndexExecutor; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexJob; -import org.openmetadata.service.apps.bundles.searchIndex.listeners.LoggingProgressListener; -import org.openmetadata.service.apps.bundles.searchIndex.listeners.QuartzProgressListener; -import org.openmetadata.service.apps.bundles.searchIndex.listeners.SlackProgressListener; import org.openmetadata.service.exception.AppException; import org.openmetadata.service.jdbi3.CollectionDAO; -import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; -import org.openmetadata.service.jdbi3.ListFilter; -import org.openmetadata.service.jdbi3.SystemRepository; -import org.openmetadata.service.search.RecreateIndexHandler; -import org.openmetadata.service.search.ReindexContext; import org.openmetadata.service.search.SearchRepository; -import org.openmetadata.service.search.vector.VectorIndexService; -import org.openmetadata.service.socket.WebSocketManager; -import org.openmetadata.service.util.FullyQualifiedName; import org.quartz.JobExecutionContext; -/** - * Quartz-scheduled application for reindexing search indices. This class handles the Quartz - * integration and delegates core reindexing logic to SearchIndexExecutor. - */ @Slf4j public class SearchIndexApp extends AbstractNativeApplication { @@ -71,10 +32,6 @@ public class SearchIndexApp extends AbstractNativeApplication { } } - private static final String ALL = "all"; - private static final String APP_SCHEDULE_RUN = "AppScheduleRun"; - private static final long WEBSOCKET_UPDATE_INTERVAL_MS = 2000; - public static final Set TIME_SERIES_ENTITIES = Set.of( ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), @@ -87,13 +44,7 @@ public class SearchIndexApp extends AbstractNativeApplication { QUERY_COST_RECORD); @Getter private EventPublisherJob jobData; - private JobExecutionContext jobExecutionContext; - private volatile boolean stopped = false; - private SearchIndexExecutor executor; - private DistributedSearchIndexExecutor distributedExecutor; - private ReindexContext recreateContext; - private RecreateIndexHandler recreateIndexHandler; - private volatile BulkSink searchIndexSink; + private volatile ReindexingOrchestrator orchestrator; public SearchIndexApp(CollectionDAO collectionDAO, SearchRepository searchRepository) { super(collectionDAO, searchRepository); @@ -105,884 +56,25 @@ public class SearchIndexApp extends AbstractNativeApplication { jobData = JsonUtils.convertValue(app.getAppConfiguration(), EventPublisherJob.class); } - private void cleanupOrphanedIndices() { - try { - OrphanedIndexCleaner cleaner = new OrphanedIndexCleaner(); - OrphanedIndexCleaner.CleanupResult result = - cleaner.cleanupOrphanedIndices(searchRepository.getSearchClient()); - if (result.deleted() > 0) { - LOG.info( - "Cleaned up {} orphaned rebuild indices on Job End (found={}, failed={})", - result.deleted(), - result.found(), - result.failed()); - } - } catch (Exception e) { - LOG.warn("Failed to cleanup orphaned indices on Job End: {}", e.getMessage()); - } - } - @Override - public void execute(JobExecutionContext jobExecutionContext) { - this.jobExecutionContext = jobExecutionContext; - initializeJobState(); - initializeJobData(jobExecutionContext); - - try { - runReindexing(jobExecutionContext); - } catch (Exception ex) { - handleExecutionException(ex); - } finally { - finalizeJobExecution(jobExecutionContext); - cleanupOrphanedIndices(); - } - } - - private void initializeJobState() { - stopped = false; - recreateContext = null; - } - - private void initializeJobData(JobExecutionContext jobExecutionContext) { - if (jobData == null) { - jobData = loadJobData(jobExecutionContext); - } - - String jobName = jobExecutionContext.getJobDetail().getKey().getName(); - if (jobName.equals(ON_DEMAND_JOB)) { - Map jsonAppConfig = - JsonUtils.convertValue(jobData, new TypeReference>() {}); - getApp().setAppConfiguration(jsonAppConfig); - } - } - - private EventPublisherJob loadJobData(JobExecutionContext jobExecutionContext) { - String appConfigJson = - (String) jobExecutionContext.getJobDetail().getJobDataMap().get(APP_CONFIG); - if (appConfigJson != null) { - return JsonUtils.readValue(appConfigJson, EventPublisherJob.class); - } - - if (getApp() != null && getApp().getAppConfiguration() != null) { - return JsonUtils.convertValue(getApp().getAppConfiguration(), EventPublisherJob.class); - } - - LOG.error("Unable to initialize jobData from JobDataMap or App configuration"); - throw new ReindexingException("JobData is not initialized"); - } - - private void cleanupOldFailures() { - try { - // Delete all previous failure records - we only keep failures for the current run - int deleted = collectionDAO.searchIndexFailureDAO().deleteAll(); - if (deleted > 0) { - LOG.info("Cleaned up {} failure records from previous runs", deleted); - } - } catch (Exception e) { - LOG.warn("Failed to cleanup old failure records", e); - } - } - - private void runReindexing(JobExecutionContext jobExecutionContext) throws Exception { - boolean success = false; - try { - if (jobData.getEntities() == null || jobData.getEntities().isEmpty()) { - LOG.info("No entities selected for reindexing, completing immediately"); - jobData.setStatus(EventPublisherJob.Status.COMPLETED); - jobData.setStats(new Stats()); - success = true; - return; - } - - setupEntities(); - cleanupOldFailures(); - - LOG.info( - "Search Index Job Started for Entities: {}, RecreateIndex: {}, DistributedIndexing: {}", - jobData.getEntities(), - jobData.getRecreateIndex(), - jobData.getUseDistributedIndexing()); - - if (Boolean.TRUE.equals(jobData.getUseDistributedIndexing())) { - runDistributedReindexing(jobExecutionContext); - success = jobData != null && jobData.getStatus() == EventPublisherJob.Status.COMPLETED; - } else { - ExecutionResult result = runSingleServerReindexing(jobExecutionContext); - success = result.isSuccessful(); - updateJobDataFromResult(result); - } - } finally { - finalizeAllEntityReindex(success); - } - } - - private ExecutionResult runSingleServerReindexing(JobExecutionContext jobExecutionContext) { - executor = new SearchIndexExecutor(collectionDAO, searchRepository); - - QuartzProgressListener quartzListener = - new QuartzProgressListener(jobExecutionContext, jobData, getApp()); - executor.addListener(quartzListener); - executor.addListener(new LoggingProgressListener()); - - if (hasSlackConfig()) { - String instanceUrl = getInstanceUrl(); - executor.addListener( - new SlackProgressListener( - jobData.getSlackBotToken(), jobData.getSlackChannel(), instanceUrl)); - } - - ReindexingJobContext context = - new QuartzJobContext( - jobExecutionContext, - getApp(), - Boolean.TRUE.equals(jobData.getUseDistributedIndexing())); - - ReindexingConfiguration config = ReindexingConfiguration.from(jobData); - - return executor.execute(config, context); - } - - private void updateJobDataFromResult(ExecutionResult result) { - if (result.finalStats() != null) { - Stats stats = result.finalStats(); - StatsReconciler.reconcile(stats); - jobData.setStats(stats); - } - - switch (result.status()) { - case COMPLETED -> jobData.setStatus(EventPublisherJob.Status.COMPLETED); - case COMPLETED_WITH_ERRORS -> jobData.setStatus(EventPublisherJob.Status.ACTIVE_ERROR); - case FAILED -> jobData.setStatus(EventPublisherJob.Status.FAILED); - case STOPPED -> jobData.setStatus(EventPublisherJob.Status.STOPPED); - } - } - - private boolean hasSlackConfig() { - return jobData.getSlackBotToken() != null - && !jobData.getSlackBotToken().isEmpty() - && jobData.getSlackChannel() != null - && !jobData.getSlackChannel().isEmpty(); - } - - private String getInstanceUrl() { - try { - SystemRepository systemRepository = Entity.getSystemRepository(); - if (systemRepository != null) { - Settings settings = systemRepository.getOMBaseUrlConfigInternal(); - if (settings != null && settings.getConfigValue() != null) { - OpenMetadataBaseUrlConfiguration urlConfig = - (OpenMetadataBaseUrlConfiguration) settings.getConfigValue(); - if (urlConfig != null && urlConfig.getOpenMetadataUrl() != null) { - return urlConfig.getOpenMetadataUrl(); - } - } - } - } catch (Exception e) { - LOG.debug("Could not get instance URL from SystemSettings", e); - } - return "http://localhost:8585"; - } - - // ========== Distributed Mode ========== - - private void runDistributedReindexing(JobExecutionContext jobExecutionContext) throws Exception { - LOG.info("Starting distributed reindexing for entities: {}", jobData.getEntities()); - - Stats stats = initializeTotalRecords(jobData.getEntities()); - jobData.setStats(stats); - - int partitionSize = jobData.getPartitionSize() != null ? jobData.getPartitionSize() : 10000; - distributedExecutor = new DistributedSearchIndexExecutor(collectionDAO, partitionSize); - distributedExecutor.performStartupRecovery(); - - // Add listeners for distributed mode (same as single-server mode) - distributedExecutor.addListener(new LoggingProgressListener()); - if (hasSlackConfig()) { - String instanceUrl = getInstanceUrl(); - distributedExecutor.addListener( - new SlackProgressListener( - jobData.getSlackBotToken(), jobData.getSlackChannel(), instanceUrl)); - } - - String createdBy = jobExecutionContext.getJobDetail().getKey().getName(); - SearchIndexJob distributedJob = - distributedExecutor.createJob(jobData.getEntities(), jobData, createdBy); - - LOG.info( - "Created distributed job {} with {} total records", - distributedJob.getId(), - distributedJob.getTotalRecords()); - - this.searchIndexSink = - searchRepository.createBulkSink( - jobData.getBatchSize(), jobData.getMaxConcurrentRequests(), jobData.getPayLoadSize()); - this.recreateIndexHandler = searchRepository.createReindexHandler(); - - if (Boolean.TRUE.equals(jobData.getRecreateIndex())) { - recreateContext = recreateIndexHandler.reCreateIndexes(jobData.getEntities()); - // Share staged index mapping with participant servers - if (recreateContext != null && !recreateContext.isEmpty()) { - distributedExecutor.updateStagedIndexMapping(recreateContext.getStagedIndexMapping()); - } - } - - updateJobStatus(EventPublisherJob.Status.RUNNING); - sendUpdates(jobExecutionContext, true); - - AppRunRecord appRecord = getJobRecord(jobExecutionContext); - distributedExecutor.setAppContext(appRecord.getAppId(), appRecord.getStartTime()); - - distributedExecutor.execute( - searchIndexSink, recreateContext, Boolean.TRUE.equals(jobData.getRecreateIndex())); - monitorDistributedJob(jobExecutionContext, distributedJob.getId()); - - if (searchIndexSink != null) { - // Wait for vector embedding tasks to complete before closing - int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); - if (pendingVectorTasks > 0) { - LOG.info("Waiting for {} pending vector embedding tasks to complete", pendingVectorTasks); - boolean vectorComplete = searchIndexSink.awaitVectorCompletion(120); - if (!vectorComplete) { - LOG.warn("Vector embedding wait timed out - some tasks may not be reflected in stats"); - } - } - - // Flush and wait for pending bulk requests - LOG.info("Flushing sink and waiting for pending bulk requests"); - boolean flushComplete = searchIndexSink.flushAndAwait(60); - if (!flushComplete) { - LOG.warn("Sink flush timed out - some requests may not be reflected in stats"); - } - - searchIndexSink.close(); - } - - SearchIndexJob finalJob = distributedExecutor.getJobWithFreshStats(); - if (finalJob != null) { - // Use actual sink stats for accurate success/failure counts - // The partition-based stats may be inaccurate because the bulk sink is asynchronous - StepStats sinkStats = searchIndexSink != null ? searchIndexSink.getStats() : null; - updateJobDataFromDistributedJob(finalJob, sinkStats); - - // Set vector stats directly from the bulk sink since the sink tracks vector - // success/failure internally and these may not be fully reflected in server stats - if (searchIndexSink != null && jobData.getStats() != null) { - StepStats sinkVectorStats = searchIndexSink.getVectorStats(); - if (sinkVectorStats != null && sinkVectorStats.getTotalRecords() > 0) { - jobData.getStats().setVectorStats(sinkVectorStats); - } - } - - saveServerStatsToJobDataMap(jobExecutionContext, finalJob); - } - - // Save stats to APP_RUN_STATS for OmAppJobListener to pick up - // This is required because distributed mode doesn't use QuartzProgressListener - if (jobData.getStats() != null) { - jobExecutionContext.getJobDetail().getJobDataMap().put(APP_RUN_STATS, jobData.getStats()); - } - - updateFinalJobStatus(); - } - - private void monitorDistributedJob( - JobExecutionContext jobExecutionContext, java.util.UUID jobId) { - CountDownLatch completionLatch = new CountDownLatch(1); - ScheduledExecutorService monitor = - Executors.newSingleThreadScheduledExecutor( - Thread.ofPlatform().name("distributed-monitor").factory()); - - try { - monitor.scheduleAtFixedRate( - () -> { - if (stopped) { - LOG.info("Stop signal received, stopping distributed job"); - distributedExecutor.stop(); - completionLatch.countDown(); - return; - } - - SearchIndexJob job = distributedExecutor.getJobWithFreshStats(); - if (job == null) { - completionLatch.countDown(); - return; - } - - IndexJobStatus status = job.getStatus(); - if (status == IndexJobStatus.COMPLETED - || status == IndexJobStatus.COMPLETED_WITH_ERRORS - || status == IndexJobStatus.FAILED - || status == IndexJobStatus.STOPPED) { - LOG.info("Distributed job {} completed with status: {}", jobId, status); - completionLatch.countDown(); - return; - } - - updateJobDataFromDistributedJob(job); - }, - 0, - WEBSOCKET_UPDATE_INTERVAL_MS, - TimeUnit.MILLISECONDS); - - completionLatch.await(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warn("Distributed job monitoring interrupted"); - } finally { - monitor.shutdownNow(); - } - } - - private void updateJobDataFromDistributedJob(SearchIndexJob distributedJob) { - updateJobDataFromDistributedJob(distributedJob, null); - } - - private void updateJobDataFromDistributedJob( - SearchIndexJob distributedJob, StepStats actualSinkStats) { - Stats stats = jobData.getStats(); - if (stats == null) { - return; - } - - // Fetch aggregated server stats once for accurate reader/sink breakdown - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats serverStatsAggr = null; - try { - serverStatsAggr = - Entity.getCollectionDAO() - .searchIndexServerStatsDAO() - .getAggregatedStats(distributedJob.getId().toString()); - if (serverStatsAggr != null) { - LOG.info( - "Fetched aggregated server stats for job {}: readerSuccess={}, readerFailed={}, " - + "sinkSuccess={}, sinkFailed={}", - distributedJob.getId(), - serverStatsAggr.readerSuccess(), - serverStatsAggr.readerFailed(), - serverStatsAggr.sinkSuccess(), - serverStatsAggr.sinkFailed()); - } - } catch (Exception e) { - LOG.debug("Could not fetch aggregated server stats for job {}", distributedJob.getId(), e); - } - - // Determine success/failed from best available source - long successRecords; - long failedRecords; - String statsSource; - - if (serverStatsAggr != null && serverStatsAggr.sinkSuccess() > 0) { - // Use server stats table (most accurate) - // processFailed = records that read successfully but failed during doc building - successRecords = serverStatsAggr.sinkSuccess(); - failedRecords = - serverStatsAggr.readerFailed() - + serverStatsAggr.sinkFailed() - + serverStatsAggr.processFailed(); - statsSource = "serverStatsTable"; - } else if (actualSinkStats != null) { - // Use local sink stats (single server scenario) - successRecords = actualSinkStats.getSuccessRecords(); - failedRecords = actualSinkStats.getFailedRecords(); - statsSource = "localSink"; - } else { - // Fallback to partition-based stats - successRecords = distributedJob.getSuccessRecords(); - failedRecords = distributedJob.getFailedRecords(); - statsSource = "partition-based"; - } - - LOG.debug( - "Stats source: {}, success={}, failed={}", statsSource, successRecords, failedRecords); - - StepStats jobStats = stats.getJobStats(); - if (jobStats != null) { - jobStats.setSuccessRecords((int) successRecords); - jobStats.setFailedRecords((int) failedRecords); - } - - StepStats readerStats = stats.getReaderStats(); - if (readerStats != null) { - readerStats.setTotalRecords((int) distributedJob.getTotalRecords()); - long readerFailed = serverStatsAggr != null ? serverStatsAggr.readerFailed() : 0; - long readerWarnings = serverStatsAggr != null ? serverStatsAggr.readerWarnings() : 0; - long readerSuccess = - serverStatsAggr != null - ? serverStatsAggr.readerSuccess() - : distributedJob.getTotalRecords() - readerFailed - readerWarnings; - readerStats.setSuccessRecords((int) readerSuccess); - readerStats.setFailedRecords((int) readerFailed); - readerStats.setWarningRecords((int) readerWarnings); - } - - // Process stats - document building stage - StepStats processStats = stats.getProcessStats(); - if (processStats != null && serverStatsAggr != null) { - long processSuccess = serverStatsAggr.processSuccess(); - long processFailed = serverStatsAggr.processFailed(); - processStats.setTotalRecords((int) (processSuccess + processFailed)); - processStats.setSuccessRecords((int) processSuccess); - processStats.setFailedRecords((int) processFailed); - } - - StepStats sinkStats = stats.getSinkStats(); - if (sinkStats != null) { - if (serverStatsAggr != null) { - // Use actual sink stats from the database - long sinkSuccess = serverStatsAggr.sinkSuccess(); - long sinkFailed = serverStatsAggr.sinkFailed(); - - // sinkTotal = docs submitted to ES = sinkSuccess + sinkFailed - long actualSinkTotal = sinkSuccess + sinkFailed; - - sinkStats.setTotalRecords((int) actualSinkTotal); - sinkStats.setSuccessRecords((int) sinkSuccess); - sinkStats.setFailedRecords((int) sinkFailed); - } else { - // Fallback: derive from reader stats (less accurate) - long readerFailed = 0; - long sinkTotal = distributedJob.getTotalRecords() - readerFailed; - sinkStats.setTotalRecords((int) sinkTotal); - sinkStats.setSuccessRecords((int) successRecords); - sinkStats.setFailedRecords((int) failedRecords); - } - } - - // Vector stats - embedding generation stage - StepStats vectorStats = stats.getVectorStats(); - if (vectorStats != null && serverStatsAggr != null) { - long vectorSuccess = serverStatsAggr.vectorSuccess(); - long vectorFailed = serverStatsAggr.vectorFailed(); - vectorStats.setTotalRecords((int) (vectorSuccess + vectorFailed)); - vectorStats.setSuccessRecords((int) vectorSuccess); - vectorStats.setFailedRecords((int) vectorFailed); - } - - if (distributedJob.getEntityStats() != null && stats.getEntityStats() != null) { - for (Map.Entry entry : - distributedJob.getEntityStats().entrySet()) { - StepStats entityStats = - stats.getEntityStats().getAdditionalProperties().get(entry.getKey()); - if (entityStats != null) { - entityStats.setSuccessRecords((int) entry.getValue().getSuccessRecords()); - entityStats.setFailedRecords((int) entry.getValue().getFailedRecords()); - } - } - } - - StatsReconciler.reconcile(stats); - - switch (distributedJob.getStatus()) { - case COMPLETED -> jobData.setStatus(EventPublisherJob.Status.COMPLETED); - case COMPLETED_WITH_ERRORS -> jobData.setStatus(EventPublisherJob.Status.ACTIVE_ERROR); - case FAILED -> jobData.setStatus(EventPublisherJob.Status.FAILED); - case STOPPING, STOPPED -> jobData.setStatus(EventPublisherJob.Status.STOPPED); - default -> jobData.setStatus(EventPublisherJob.Status.RUNNING); - } - } - - private void saveServerStatsToJobDataMap( - JobExecutionContext jobExecutionContext, SearchIndexJob distributedJob) { - try { - AppRunRecord appRecord = getJobRecord(jobExecutionContext); - SuccessContext successContext = appRecord.getSuccessContext(); - if (successContext == null) { - successContext = new SuccessContext(); - } - - if (distributedJob.getServerStats() != null && !distributedJob.getServerStats().isEmpty()) { - LOG.info( - "Saving serverStats to job data map: {} servers with data: {}", - distributedJob.getServerStats().size(), - distributedJob.getServerStats()); - successContext.withAdditionalProperty("serverStats", distributedJob.getServerStats()); - successContext.withAdditionalProperty( - "serverCount", distributedJob.getServerStats().size()); - successContext.withAdditionalProperty( - "distributedJobId", distributedJob.getId().toString()); - } else { - LOG.warn( - "No server stats available for distributed job {} - serverStats is {} ", - distributedJob.getId(), - distributedJob.getServerStats() == null ? "null" : "empty"); - } - - if (jobData.getStats() != null) { - successContext.withAdditionalProperty("stats", jobData.getStats()); - } - - appRecord.setSuccessContext(successContext); - jobExecutionContext - .getJobDetail() - .getJobDataMap() - .put("AppScheduleRun", JsonUtils.pojoToJson(appRecord)); - - } catch (Exception e) { - LOG.error("Failed to save serverStats to job data map", e); - } - } - - // ========== Helper Methods ========== - - private void setupEntities() { - boolean containsAll = jobData.getEntities().contains(ALL); - if (containsAll) { - jobData.setEntities(getAll()); - } - } - - private Set getAll() { - Set entities = - new HashSet<>( - Entity.getEntityList().stream() - .filter(t -> searchRepository.getEntityIndexMap().containsKey(t)) - .toList()); - entities.addAll( - TIME_SERIES_ENTITIES.stream() - .filter(t -> searchRepository.getEntityIndexMap().containsKey(t)) - .toList()); - return entities; - } - - public Stats initializeTotalRecords(Set entities) { - Stats stats = new Stats(); - stats.setEntityStats(new org.openmetadata.schema.system.EntityStats()); - stats.setJobStats(new StepStats()); - stats.setReaderStats(new StepStats()); - stats.setProcessStats(new StepStats()); - stats.setSinkStats(new StepStats()); - stats.setVectorStats(new StepStats()); - - int total = 0; - for (String entityType : entities) { - int entityTotal = getEntityTotal(entityType); - total += entityTotal; - - StepStats entityStats = new StepStats(); - entityStats.setTotalRecords(entityTotal); - entityStats.setSuccessRecords(0); - entityStats.setFailedRecords(0); - stats.getEntityStats().getAdditionalProperties().put(entityType, entityStats); - } - - stats.getJobStats().setTotalRecords(total); - stats.getJobStats().setSuccessRecords(0); - stats.getJobStats().setFailedRecords(0); - - stats.getReaderStats().setTotalRecords(total); - stats.getReaderStats().setSuccessRecords(0); - stats.getReaderStats().setFailedRecords(0); - - stats.getProcessStats().setTotalRecords(0); - stats.getProcessStats().setSuccessRecords(0); - stats.getProcessStats().setFailedRecords(0); - - stats.getSinkStats().setTotalRecords(0); - stats.getSinkStats().setSuccessRecords(0); - stats.getSinkStats().setFailedRecords(0); - - stats.getVectorStats().setTotalRecords(0); - stats.getVectorStats().setSuccessRecords(0); - stats.getVectorStats().setFailedRecords(0); - - return stats; - } - - private int getEntityTotal(String entityType) { - try { - String correctedType = "queryCostResult".equals(entityType) ? QUERY_COST_RECORD : entityType; - - if (!TIME_SERIES_ENTITIES.contains(correctedType)) { - return Entity.getEntityRepository(correctedType).getDao().listTotalCount(); - } else { - ListFilter listFilter = new ListFilter(null); - EntityTimeSeriesRepository repository; - - if (isDataInsightIndex(correctedType)) { - listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(correctedType)); - repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); - } else { - repository = Entity.getEntityTimeSeriesRepository(correctedType); - } - - return repository.getTimeSeriesDao().listCount(listFilter); - } - } catch (Exception e) { - LOG.debug("Error getting total for '{}'", entityType, e); - return 0; - } - } - - private boolean isDataInsightIndex(String entityType) { - return entityType.endsWith("ReportData"); - } - - private void updateJobStatus(EventPublisherJob.Status newStatus) { - if (stopped - && newStatus != EventPublisherJob.Status.STOP_IN_PROGRESS - && newStatus != EventPublisherJob.Status.STOPPED) { - return; - } - jobData.setStatus(newStatus); - } - - private void updateFinalJobStatus() { - if (stopped) { - updateJobStatus(EventPublisherJob.Status.STOPPED); - } else if (hasIncompleteProcessing()) { - updateJobStatus(EventPublisherJob.Status.ACTIVE_ERROR); - } else { - updateJobStatus(EventPublisherJob.Status.COMPLETED); - } - } - - private boolean hasIncompleteProcessing() { - if (jobData == null || jobData.getStats() == null || jobData.getStats().getJobStats() == null) { - return false; - } - - StepStats jobStats = jobData.getStats().getJobStats(); - long failed = jobStats.getFailedRecords() != null ? jobStats.getFailedRecords() : 0; - long processed = jobStats.getSuccessRecords() != null ? jobStats.getSuccessRecords() : 0; - long total = jobStats.getTotalRecords() != null ? jobStats.getTotalRecords() : 0; - - return failed > 0 || (total > 0 && processed < total); - } - - private void finalizeAllEntityReindex(boolean finalSuccess) { - if (recreateIndexHandler == null || recreateContext == null) { - return; - } - - // Get already-promoted entities from distributed executor (if running in distributed mode) - Set promotedEntities = Collections.emptySet(); - if (distributedExecutor != null && distributedExecutor.getEntityTracker() != null) { - promotedEntities = distributedExecutor.getEntityTracker().getPromotedEntities(); - } - - // Calculate entities that still need finalization - Set entitiesToFinalize = new HashSet<>(recreateContext.getEntities()); - entitiesToFinalize.removeAll(promotedEntities); - - // Vector index is a pseudo-entity with no partitions or batch tracking — handle separately - boolean hasVectorIndex = entitiesToFinalize.remove(VectorIndexService.VECTOR_INDEX_KEY); - - try { - if (!entitiesToFinalize.isEmpty()) { - LOG.info( - "Finalizing {} remaining entities (already promoted: {})", - entitiesToFinalize.size(), - promotedEntities.size()); - - for (String entityType : entitiesToFinalize) { - try { - finalizeEntityReindex(entityType, finalSuccess); - } catch (Exception ex) { - LOG.error("Failed to finalize reindex for entity: {}", entityType, ex); - } - } - } - - if (hasVectorIndex) { - finalizeVectorIndex(finalSuccess); - } - } finally { - recreateContext = null; - } - } - - private void finalizeVectorIndex(boolean finalSuccess) { - // Vector index data is written as a side-effect of processing real entities. - // Promote when the job ran to completion (even with some errors) since partial - // vector data is better than an orphaned rebuild index. Only discard on - // FAILED (job crashed) or STOPPED (user cancelled). - boolean vectorSuccess = - finalSuccess - || (jobData != null && jobData.getStatus() == EventPublisherJob.Status.ACTIVE_ERROR); - - try { - finalizeEntityReindex(VectorIndexService.VECTOR_INDEX_KEY, vectorSuccess); - } catch (Exception ex) { - LOG.error("Failed to finalize vector index", ex); - } - } - - private void finalizeEntityReindex(String entityType, boolean success) { - if (recreateIndexHandler == null || recreateContext == null) { - return; - } - - try { - var entityReindexContext = - org.openmetadata.service.search.EntityReindexContext.builder() - .entityType(entityType) - .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) - .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) - .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) - .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) - .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) - .existingAliases(recreateContext.getExistingAliases(entityType)) - .parentAliases( - new HashSet<>(listOrEmpty(recreateContext.getParentAliases(entityType)))) - .build(); - - recreateIndexHandler.finalizeReindex(entityReindexContext, success); - } catch (Exception ex) { - LOG.error("Failed to finalize index recreation flow", ex); - } - } - - private void handleExecutionException(Exception ex) { - BulkSink sink = searchIndexSink; - if (sink != null) { - searchIndexSink = null; - try { - sink.close(); - } catch (Exception e) { - LOG.error("Error closing search index sink", e); - } - } - - if (executor != null && jobData != null) { - try { - Stats executorStats = executor.getStats().get(); - if (executorStats != null) { - jobData.setStats(executorStats); - } - } catch (Exception e) { - LOG.debug("Could not capture executor stats during exception handling", e); - } - } - - if (stopped) { - if (jobData != null) { - jobData.setStatus(EventPublisherJob.Status.STOPPED); - } - } else { - IndexingError error = - new IndexingError() - .withErrorSource(IndexingError.ErrorSource.JOB) - .withMessage("Reindexing Job Exception: " + ex.getMessage()); - LOG.error("Reindexing Job Failed", ex); - - if (jobData != null) { - jobData.setStatus(EventPublisherJob.Status.FAILED); - jobData.setFailure(error); - } - } - } - - private void finalizeJobExecution(JobExecutionContext jobExecutionContext) { - sendUpdates(jobExecutionContext, true); - - if (stopped && jobExecutionContext != null) { - AppRunRecord appRecord = getJobRecord(jobExecutionContext); - appRecord.setStatus(AppRunRecord.Status.STOPPED); - jobExecutionContext - .getJobDetail() - .getJobDataMap() - .put(APP_SCHEDULE_RUN, JsonUtils.pojoToJson(appRecord)); - } - } - - private void sendUpdates(JobExecutionContext jobExecutionContext, boolean force) { - try { - updateRecordToDbAndNotify(jobExecutionContext); - } catch (Exception ex) { - LOG.error("Failed to send updates", ex); - } - } - - public void updateRecordToDbAndNotify(JobExecutionContext jobExecutionContext) { - AppRunRecord appRecord = getJobRecord(jobExecutionContext); - appRecord.setStatus(AppRunRecord.Status.fromValue(jobData.getStatus().value())); - - if (jobData.getFailure() != null) { - appRecord.setFailureContext( - new FailureContext().withAdditionalProperty("failure", jobData.getFailure())); - } - - if (jobData.getStats() != null) { - SuccessContext successContext = - new SuccessContext().withAdditionalProperty("stats", jobData.getStats()); - - SearchIndexJob distributedJob = - distributedExecutor != null ? distributedExecutor.getJobWithFreshStats() : null; - - try { - String jobIdStr = - distributedJob != null - ? distributedJob.getId().toString() - : getApp().getId().toString(); - int failureCount = collectionDAO.searchIndexFailureDAO().countByJobId(jobIdStr); - if (failureCount > 0) { - successContext.withAdditionalProperty("failureRecordCount", failureCount); - } - } catch (Exception e) { - LOG.debug("Could not get failure count", e); - } - - if (distributedJob != null && distributedJob.getServerStats() != null) { - successContext.withAdditionalProperty("serverStats", distributedJob.getServerStats()); - successContext.withAdditionalProperty( - "serverCount", distributedJob.getServerStats().size()); - successContext.withAdditionalProperty( - "distributedJobId", distributedJob.getId().toString()); - } - - appRecord.setSuccessContext(successContext); - } - - if (WebSocketManager.getInstance() != null) { - String messageJson = JsonUtils.pojoToJson(appRecord); - WebSocketManager.getInstance() - .broadCastMessageToAll(SEARCH_INDEX_JOB_BROADCAST_CHANNEL, messageJson); - } + public void execute(JobExecutionContext ctx) { + OrchestratorContext orchCtx = + new QuartzOrchestratorContext( + ctx, getApp(), this::getJobRecord, this::pushAppStatusUpdates); + ReindexingOrchestrator orch = + new ReindexingOrchestrator(collectionDAO, searchRepository, orchCtx); + this.orchestrator = orch; + orch.run(jobData); + this.jobData = orch.getJobData(); } @Override public void stop() { - LOG.info("Reindexing job is being stopped."); - stopped = true; - - if (executor != null) { - executor.stop(); + ReindexingOrchestrator orch = this.orchestrator; + if (orch != null) { + orch.stop(); + this.jobData = orch.getJobData(); } - - if (distributedExecutor != null) { - try { - distributedExecutor.stop(); - } catch (Exception e) { - LOG.error("Error stopping distributed executor", e); - } - } - - if (jobData != null) { - jobData.setStatus(EventPublisherJob.Status.STOPPED); - } - - if (jobExecutionContext != null) { - AppRunRecord appRecord = getJobRecord(jobExecutionContext); - appRecord.setStatus(AppRunRecord.Status.STOPPED); - appRecord.setEndTime(System.currentTimeMillis()); - jobExecutionContext - .getJobDetail() - .getJobDataMap() - .put(APP_SCHEDULE_RUN, JsonUtils.pojoToJson(appRecord)); - pushAppStatusUpdates(jobExecutionContext, appRecord, true); - sendUpdates(jobExecutionContext, true); - } - - BulkSink sink = searchIndexSink; - if (sink != null) { - searchIndexSink = null; - try { - sink.close(); - } catch (Exception e) { - LOG.error("Error closing search index sink", e); - } - } - - LOG.info("Reindexing job stopped successfully."); } @Override diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java index bcfe3a26358..4f6ad4c55f3 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java @@ -24,7 +24,9 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -45,6 +47,7 @@ import org.openmetadata.schema.utils.ResultList; import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.stats.EntityStatsTracker; import org.openmetadata.service.apps.bundles.searchIndex.stats.JobStatsManager; +import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; import org.openmetadata.service.exception.SearchIndexException; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.jdbi3.EntityRepository; @@ -54,13 +57,13 @@ import org.openmetadata.service.search.DefaultRecreateHandler; import org.openmetadata.service.search.EntityReindexContext; import org.openmetadata.service.search.RecreateIndexHandler; import org.openmetadata.service.search.ReindexContext; -import org.openmetadata.service.search.SearchClusterMetrics; import org.openmetadata.service.search.SearchRepository; import org.openmetadata.service.util.FullyQualifiedName; import org.openmetadata.service.util.RestUtil; import org.openmetadata.service.workflows.interfaces.Source; import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; +import org.slf4j.MDC; /** * Core reindexing executor that handles entity indexing without any Quartz dependencies. Can be @@ -142,6 +145,7 @@ public class SearchIndexExecutor implements AutoCloseable { private final Map entityBatchCounters = new ConcurrentHashMap<>(); private final Map entityBatchFailures = new ConcurrentHashMap<>(); private final Set promotedEntities = ConcurrentHashMap.newKeySet(); + private final Map sinkTrackers = new ConcurrentHashMap<>(); record IndexingTask(String entityType, ResultList entities, int offset, int retryCount) { IndexingTask(String entityType, ResultList entities, int offset) { @@ -258,14 +262,16 @@ public class SearchIndexExecutor implements AutoCloseable { entityBatchCounters.clear(); entityBatchFailures.clear(); promotedEntities.clear(); + sinkTrackers.clear(); initStatsManager(); } private ExecutionResult executeSingleServer() throws Exception { Set entities = expandEntities(config.entities()); - ReindexingConfiguration effectiveConfig = applyAutoTuning(entities); + batchSize.set(config.batchSize()); + originalBatchSize.set(config.batchSize()); - listeners.onJobConfigured(context, effectiveConfig); + listeners.onJobConfigured(context, config); stats.set(initializeTotalRecords(entities)); @@ -284,20 +290,15 @@ public class SearchIndexExecutor implements AutoCloseable { this.failureRecorder = new IndexingFailureRecorder(collectionDAO, jobId, serverId); cleanupOldFailures(); - initializeSink(effectiveConfig); + initializeSink(config); - if (effectiveConfig.recreateIndex()) { + if (config.recreateIndex()) { validateClusterCapacity(entities); listeners.onIndexRecreationStarted(entities); recreateContext = reCreateIndexes(entities); } - SearchClusterMetrics clusterMetrics = null; - if (effectiveConfig.autoTune()) { - clusterMetrics = fetchClusterMetrics(); - } - - reIndexFromStartToEnd(clusterMetrics, entities); + reIndexFromStartToEnd(entities); closeSinkIfNeeded(); // Promote anything yet to be promoted such as vector search indexes which is not part of // entities set @@ -313,59 +314,6 @@ public class SearchIndexExecutor implements AutoCloseable { return entities; } - private ReindexingConfiguration applyAutoTuning(Set entities) { - if (!config.autoTune()) { - batchSize.set(config.batchSize()); - originalBatchSize.set(config.batchSize()); - return config; - } - - SearchClusterMetrics metrics = fetchClusterMetrics(); - if (metrics == null) { - batchSize.set(config.batchSize()); - originalBatchSize.set(config.batchSize()); - return config; - } - - batchSize.set(metrics.getRecommendedBatchSize()); - originalBatchSize.set(metrics.getRecommendedBatchSize()); - - return ReindexingConfiguration.builder() - .entities(entities) - .batchSize(metrics.getRecommendedBatchSize()) - .consumerThreads(metrics.getRecommendedConsumerThreads()) - .producerThreads(metrics.getRecommendedProducerThreads()) - .queueSize(metrics.getRecommendedQueueSize()) - .maxConcurrentRequests(metrics.getRecommendedConcurrentRequests()) - .payloadSize(metrics.getMaxPayloadSizeBytes()) - .recreateIndex(config.recreateIndex()) - .autoTune(true) - .useDistributedIndexing(config.useDistributedIndexing()) - .force(config.force()) - .maxRetries(config.maxRetries()) - .initialBackoff(config.initialBackoff()) - .maxBackoff(config.maxBackoff()) - .searchIndexMappingLanguage(config.searchIndexMappingLanguage()) - .afterCursor(config.afterCursor()) - .slackBotToken(config.slackBotToken()) - .slackChannel(config.slackChannel()) - .build(); - } - - private SearchClusterMetrics fetchClusterMetrics() { - try { - long totalRecords = - stats.get() != null && stats.get().getJobStats() != null - ? stats.get().getJobStats().getTotalRecords() - : 0; - return SearchClusterMetrics.fetchClusterMetrics( - searchRepository, totalRecords, searchRepository.getMaxDBConnections()); - } catch (Exception e) { - LOG.warn("Failed to fetch cluster metrics, using defaults", e); - return null; - } - } - private void validateClusterCapacity(Set entities) { try { SearchIndexClusterValidator validator = new SearchIndexClusterValidator(); @@ -392,9 +340,17 @@ public class SearchIndexExecutor implements AutoCloseable { } private void handleSinkFailure( - String entityType, String entityId, String entityFqn, String errorMessage) { + String entityType, + String entityId, + String entityFqn, + String errorMessage, + IndexingFailureRecorder.FailureStage stage) { if (failureRecorder != null) { - failureRecorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + if (stage == IndexingFailureRecorder.FailureStage.PROCESS) { + failureRecorder.recordProcessFailure(entityType, entityId, entityFqn, errorMessage); + } else { + failureRecorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + } } } @@ -410,14 +366,13 @@ public class SearchIndexExecutor implements AutoCloseable { } } - private void reIndexFromStartToEnd(SearchClusterMetrics clusterMetrics, Set entities) - throws InterruptedException { + private void reIndexFromStartToEnd(Set entities) throws InterruptedException { long totalEntities = stats.get() != null && stats.get().getJobStats() != null ? stats.get().getJobStats().getTotalRecords() : 0; - ThreadConfiguration threadConfig = calculateThreadConfiguration(totalEntities, clusterMetrics); + ThreadConfiguration threadConfig = calculateThreadConfiguration(totalEntities); int effectiveQueueSize = initializeQueueAndExecutors(threadConfig, entities.size()); LOG.info( @@ -429,16 +384,13 @@ public class SearchIndexExecutor implements AutoCloseable { executeReindexing(threadConfig.numConsumers(), entities); } - private ThreadConfiguration calculateThreadConfiguration( - long totalEntities, SearchClusterMetrics clusterMetrics) { + private ThreadConfiguration calculateThreadConfiguration(long totalEntities) { int numConsumers = config.consumerThreads() > 0 ? Math.min(config.consumerThreads(), MAX_CONSUMER_THREADS) : 2; - int numProducers = Math.clamp((int) (totalEntities / 10000), 2, MAX_PRODUCER_THREADS); - - if (clusterMetrics != null) { - numConsumers = Math.min(clusterMetrics.getRecommendedConsumerThreads(), MAX_CONSUMER_THREADS); - numProducers = Math.min(clusterMetrics.getRecommendedProducerThreads(), MAX_PRODUCER_THREADS); - } + int numProducers = + config.producerThreads() > 1 + ? Math.min(config.producerThreads(), MAX_PRODUCER_THREADS) + : Math.clamp((int) (totalEntities / 10000), 2, MAX_PRODUCER_THREADS); return adjustThreadsForLimit(numProducers, numConsumers); } @@ -465,8 +417,12 @@ public class SearchIndexExecutor implements AutoCloseable { taskQueue = new LinkedBlockingQueue<>(effectiveQueueSize); producersDone.set(false); + int maxJobThreads = + Math.max(1, MAX_TOTAL_THREADS - threadConfig.numProducers() - threadConfig.numConsumers()); + int cappedEntityCount = Math.min(entityCount, maxJobThreads); jobExecutor = - Executors.newFixedThreadPool(entityCount, Thread.ofPlatform().name("job-", 0).factory()); + Executors.newFixedThreadPool( + cappedEntityCount, Thread.ofPlatform().name("job-", 0).factory()); int finalNumConsumers = Math.min(threadConfig.numConsumers(), MAX_CONSUMER_THREADS); consumerExecutor = @@ -484,7 +440,8 @@ public class SearchIndexExecutor implements AutoCloseable { MemoryInfo memInfo = new MemoryInfo(); long estimatedEntitySize = 5 * 1024L; long maxQueueMemory = (long) (memInfo.maxMemory * 0.25); - int memoryBasedLimit = (int) (maxQueueMemory / (estimatedEntitySize * batchSize.get())); + long memoryBasedLimitLong = maxQueueMemory / (estimatedEntitySize * batchSize.get()); + int memoryBasedLimit = (int) Math.max(1, Math.min(memoryBasedLimitLong, Integer.MAX_VALUE)); return Math.min(requestedSize, memoryBasedLimit); } @@ -508,9 +465,18 @@ public class SearchIndexExecutor implements AutoCloseable { private CountDownLatch startConsumerThreads(int numConsumers) { CountDownLatch consumerLatch = new CountDownLatch(numConsumers); + Map mdc = MDC.getCopyOfContextMap(); for (int i = 0; i < numConsumers; i++) { final int consumerId = i; - consumerExecutor.submit(() -> runConsumer(consumerId, consumerLatch)); + consumerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + runConsumer(consumerId, consumerLatch); + } finally { + MDC.clear(); + } + }); } return consumerLatch; } @@ -520,7 +486,7 @@ public class SearchIndexExecutor implements AutoCloseable { try { while (!stopped.get()) { try { - IndexingTask task = taskQueue.poll(500, TimeUnit.MILLISECONDS); + IndexingTask task = taskQueue.poll(200, TimeUnit.MILLISECONDS); if (task == null) { continue; } @@ -594,12 +560,35 @@ public class SearchIndexExecutor implements AutoCloseable { contextData.put(ENTITY_TYPE_KEY, entityType); contextData.put(RECREATE_INDEX, config.recreateIndex()); contextData.put(RECREATE_CONTEXT, recreateContext); - contextData.put(BulkSink.STATS_TRACKER_CONTEXT_KEY, getTracker(entityType)); + contextData.put(BulkSink.STATS_TRACKER_CONTEXT_KEY, getSinkTracker(entityType)); getTargetIndexForEntity(entityType) .ifPresent(index -> contextData.put(TARGET_INDEX_KEY, index)); return contextData; } + private StageStatsTracker getSinkTracker(String entityType) { + if (context == null) { + return null; + } + return sinkTrackers.computeIfAbsent( + entityType, + et -> { + String jobId = context.getJobId().toString(); + String serverId = + org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .ServerIdentityResolver + .getInstance() + .getServerId(); + return new StageStatsTracker( + jobId, serverId, et, collectionDAO.searchIndexServerStatsDAO()); + }); + } + private void writeEntitiesToSink( String entityType, ResultList entities, Map contextData) throws Exception { if (!TIME_SERIES_ENTITIES.contains(entityType)) { @@ -700,6 +689,12 @@ public class SearchIndexExecutor implements AutoCloseable { if (errorMessage != null && isBackpressureError(errorMessage)) { consecutiveErrors.incrementAndGet(); consecutiveSuccesses.set(0); + + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordBackpressureEvent(); + } + LOG.warn("Detected backpressure (consecutive errors: {})", consecutiveErrors.get()); boolean isPayloadTooLarge = isPayloadTooLargeError(errorMessage); @@ -768,7 +763,26 @@ public class SearchIndexExecutor implements AutoCloseable { new TuningContext( new MemoryInfo(), batchSize.get(), consecutiveErrors.get(), consecutiveSuccesses.get()); - if (tuningContext.errorCount == 0 + // Critical memory tier (>90%): halve batch size aggressively + if (tuningContext.memInfo.usageRatio > 0.9) { + int newBatchSize = Math.max(tuningContext.currentBatchSize / 2, 25); + if (newBatchSize != tuningContext.currentBatchSize) { + batchSize.set(newBatchSize); + LOG.warn( + "Auto-tune: Aggressively reduced batch size to {} due to critical memory ({}% used)", + newBatchSize, (int) (tuningContext.memInfo.usageRatio * 100)); + updateSinkBatchSize(newBatchSize); + } + } else if (tuningContext.memInfo.usageRatio > 0.8) { + int newBatchSize = Math.max(tuningContext.currentBatchSize - 100, 50); + if (newBatchSize != tuningContext.currentBatchSize) { + batchSize.set(newBatchSize); + LOG.warn( + "Auto-tune: Reduced batch size to {} due to memory pressure ({}% used)", + newBatchSize, (int) (tuningContext.memInfo.usageRatio * 100)); + updateSinkBatchSize(newBatchSize); + } + } else if (tuningContext.errorCount == 0 && tuningContext.successCount > BATCH_SIZE_INCREASE_THRESHOLD && tuningContext.memInfo.usageRatio < 0.7) { int newBatchSize = Math.min(tuningContext.currentBatchSize + 50, 1000); @@ -780,14 +794,6 @@ public class SearchIndexExecutor implements AutoCloseable { String.format("%.1f", currentThroughput)); updateSinkBatchSize(newBatchSize); } - } else if (tuningContext.memInfo.usageRatio > 0.8) { - int newBatchSize = Math.max(tuningContext.currentBatchSize - 100, 50); - if (newBatchSize != tuningContext.currentBatchSize) { - batchSize.set(newBatchSize); - LOG.warn( - "Auto-tune: Reduced batch size to {} due to memory pressure ({}% used)", - newBatchSize, (int) (tuningContext.memInfo.usageRatio * 100)); - } } } @@ -799,13 +805,10 @@ public class SearchIndexExecutor implements AutoCloseable { } } - private void signalConsumersToStop(int numConsumers) { + private void signalConsumersToStop(int numConsumers) throws InterruptedException { producersDone.set(true); for (int i = 0; i < numConsumers; i++) { - boolean offered = taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); - if (!offered) { - LOG.debug("Could not add poison pill to queue"); - } + taskQueue.put(new IndexingTask<>(POISON_PILL, null, -1)); } } @@ -817,27 +820,47 @@ public class SearchIndexExecutor implements AutoCloseable { } private void processEntityReindex(Set entities) throws InterruptedException { - int snapshotBatchSize = batchSize.get(); - int latchCount = getTotalLatchCount(entities, snapshotBatchSize); - CountDownLatch producerLatch = new CountDownLatch(latchCount); + // Use Phaser instead of pre-computed CountDownLatch to handle dynamic reader counts. + // Each entity type registers as a party, then dynamically registers its actual readers. + // This eliminates the batch-size-snapshot mismatch where auto-tune could desynchronize + // the pre-computed latch count from the actual number of readers created. + List ordered = EntityPriority.sortByPriority(entities); + LOG.info("Entity processing order: {}", ordered); + Phaser producerPhaser = new Phaser(entities.size()); + Map mdc = MDC.getCopyOfContextMap(); - for (String entityType : entities) { - jobExecutor.submit(() -> processEntityType(entityType, producerLatch, snapshotBatchSize)); + for (String entityType : ordered) { + jobExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + processEntityType(entityType, producerPhaser); + } finally { + MDC.clear(); + } + }); } - while (!producerLatch.await(1, TimeUnit.SECONDS)) { + int phase = 0; + while (!producerPhaser.isTerminated()) { if (stopped.get() || Thread.currentThread().isInterrupted()) { LOG.info("Stop signal received during reindexing"); if (producerExecutor != null) producerExecutor.shutdownNow(); if (jobExecutor != null) jobExecutor.shutdownNow(); return; } + try { + producerPhaser.awaitAdvanceInterruptibly(phase, 1, TimeUnit.SECONDS); + break; + } catch (TimeoutException e) { + // Continue checking stop signal + } } } - private void processEntityType( - String entityType, CountDownLatch producerLatch, int fixedBatchSize) { + private void processEntityType(String entityType, Phaser producerPhaser) { try { + int fixedBatchSize = EntityBatchSizeEstimator.estimateBatchSize(entityType, batchSize.get()); int totalEntityRecords = getTotalEntityRecords(entityType); listeners.onEntityTypeStarted(entityType, totalEntityRecords); @@ -850,55 +873,91 @@ public class SearchIndexExecutor implements AutoCloseable { MAX_READERS_PER_ENTITY); entityBatchCounters.put(entityType, new AtomicInteger(numReaders)); - if (TIME_SERIES_ENTITIES.contains(entityType)) { - submitReaders( + // Dynamically register actual readers with the phaser + producerPhaser.bulkRegister(numReaders); + + try { + if (TIME_SERIES_ENTITIES.contains(entityType)) { + Long filterStartTs = null; + Long filterEndTs = null; + if (config != null) { + long startTs = config.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + filterStartTs = startTs; + filterEndTs = System.currentTimeMillis(); + } + } + final Long tsStart = filterStartTs; + final Long tsEnd = filterEndTs; + submitReaders( + entityType, + totalEntityRecords, + fixedBatchSize, + numReaders, + producerPhaser, + () -> { + PaginatedEntityTimeSeriesSource source = + (tsStart != null) + ? new PaginatedEntityTimeSeriesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords, + tsStart, + tsEnd) + : new PaginatedEntityTimeSeriesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords); + return source::readWithCursor; + }, + (readers, total) -> { + List cursors = new ArrayList<>(); + int perReader = total / readers; + for (int i = 1; i < readers; i++) { + cursors.add(RestUtil.encodeCursor(String.valueOf(i * perReader))); + } + return cursors; + }); + } else { + PaginatedEntitiesSource entSource = + new PaginatedEntitiesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords); + submitReaders( + entityType, + totalEntityRecords, + fixedBatchSize, + numReaders, + producerPhaser, + () -> { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords); + return source::readNextKeyset; + }, + entSource::findBoundaryCursors); + } + } catch (Exception e) { + LOG.error( + "Failed to submit readers for {}, deregistering {} phaser parties", entityType, - totalEntityRecords, - fixedBatchSize, numReaders, - producerLatch, - () -> { - PaginatedEntityTimeSeriesSource source = - new PaginatedEntityTimeSeriesSource( - entityType, - fixedBatchSize, - getSearchIndexFields(entityType), - totalEntityRecords); - return source::readWithCursor; - }, - (readers, total) -> { - List cursors = new ArrayList<>(); - int perReader = total / readers; - for (int i = 1; i < readers; i++) { - cursors.add(RestUtil.encodeCursor(String.valueOf(i * perReader))); - } - return cursors; - }); - } else { - PaginatedEntitiesSource entSource = - new PaginatedEntitiesSource( - entityType, fixedBatchSize, getSearchIndexFields(entityType), totalEntityRecords); - submitReaders( - entityType, - totalEntityRecords, - fixedBatchSize, - numReaders, - producerLatch, - () -> { - PaginatedEntitiesSource source = - new PaginatedEntitiesSource( - entityType, - fixedBatchSize, - getSearchIndexFields(entityType), - totalEntityRecords); - return source::readNextKeyset; - }, - entSource::findBoundaryCursors); + e); + for (int i = 0; i < numReaders; i++) { + producerPhaser.arriveAndDeregister(); + } + throw e; } } else { entityBatchCounters.put(entityType, new AtomicInteger(1)); promoteEntityIndexIfReady(entityType); - producerLatch.countDown(); } StepStats entityStats = @@ -908,6 +967,9 @@ public class SearchIndexExecutor implements AutoCloseable { listeners.onEntityTypeCompleted(entityType, entityStats); } catch (Exception e) { LOG.error("Error processing entity type {}", entityType, e); + } finally { + // Deregister the entity coordinator party + producerPhaser.arriveAndDeregister(); } } @@ -916,21 +978,29 @@ public class SearchIndexExecutor implements AutoCloseable { int totalRecords, int fixedBatchSize, int numReaders, - CountDownLatch producerLatch, + Phaser producerPhaser, java.util.function.Supplier readerFactory, java.util.function.BiFunction> boundaryFinder) { + Map mdc = MDC.getCopyOfContextMap(); if (numReaders == 1) { KeysetBatchReader reader = readerFactory.get(); producerExecutor.submit( - () -> + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { processKeysetBatches( - entityType, totalRecords, fixedBatchSize, null, reader, producerLatch)); + entityType, Integer.MAX_VALUE, fixedBatchSize, null, reader, producerPhaser); + } finally { + MDC.clear(); + } + }); return; } List boundaries = boundaryFinder.apply(numReaders, totalRecords); int actualReaders = boundaries.size() + 1; - int recordsPerReader = totalRecords / actualReaders; + // Use ceiling division to avoid rounding-related entity loss at reader boundaries + int recordsPerReader = (totalRecords + actualReaders - 1) / actualReaders; if (actualReaders < numReaders) { LOG.warn( @@ -940,28 +1010,70 @@ public class SearchIndexExecutor implements AutoCloseable { numReaders - 1, actualReaders); entityBatchCounters.get(entityType).set(actualReaders); + // Deregister extra reader parties from the phaser for (int j = 0; j < numReaders - actualReaders; j++) { - producerLatch.countDown(); + producerPhaser.arriveAndDeregister(); } } for (int i = 0; i < actualReaders; i++) { String startCursor = (i == 0) ? null : boundaries.get(i - 1); - int limit = - (i == actualReaders - 1) - ? totalRecords - recordsPerReader * (actualReaders - 1) - : recordsPerReader; + String endCursorForReader = (i < boundaries.size()) ? boundaries.get(i) : null; + int limit = (i == actualReaders - 1) ? Integer.MAX_VALUE : recordsPerReader; KeysetBatchReader readerSource = readerFactory.get(); final int readerLimit = limit; + final String readerEndCursor = endCursorForReader; producerExecutor.submit( - () -> + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { processKeysetBatches( entityType, readerLimit, fixedBatchSize, startCursor, readerSource, - producerLatch)); + producerPhaser, + readerEndCursor); + } finally { + MDC.clear(); + } + }); + } + } + + private boolean hasReachedEndCursor(String afterCursor, String endCursor) { + if (endCursor == null || afterCursor == null) return false; + String decodedAfter = RestUtil.decodeCursor(afterCursor); + String decodedEnd = RestUtil.decodeCursor(endCursor); + if (decodedAfter == null || decodedEnd == null) return false; + + // Time-series cursors are numeric offsets + try { + int afterOffset = Integer.parseInt(decodedAfter); + int endOffset = Integer.parseInt(decodedEnd); + return afterOffset >= endOffset; + } catch (NumberFormatException ignored) { + // Not a numeric cursor, fall through to JSON comparison + } + + // Regular entity cursors are JSON maps with "name" and "id" fields + try { + @SuppressWarnings("unchecked") + Map afterMap = + org.openmetadata.schema.utils.JsonUtils.readValue(decodedAfter, Map.class); + @SuppressWarnings("unchecked") + Map endMap = + org.openmetadata.schema.utils.JsonUtils.readValue(decodedEnd, Map.class); + String afterName = afterMap.getOrDefault("name", ""); + String endName = endMap.getOrDefault("name", ""); + int nameCompare = afterName.compareTo(endName); + if (nameCompare != 0) return nameCompare >= 0; + String afterId = afterMap.getOrDefault("id", ""); + String endId = endMap.getOrDefault("id", ""); + return afterId.compareTo(endId) >= 0; + } catch (Exception e) { + return decodedAfter.compareTo(decodedEnd) >= 0; } } @@ -971,7 +1083,19 @@ public class SearchIndexExecutor implements AutoCloseable { int fixedBatchSize, String startCursor, KeysetBatchReader batchReader, - CountDownLatch producerLatch) { + Phaser producerPhaser) { + processKeysetBatches( + entityType, recordLimit, fixedBatchSize, startCursor, batchReader, producerPhaser, null); + } + + private void processKeysetBatches( + String entityType, + int recordLimit, + int fixedBatchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser producerPhaser, + String endCursor) { boolean hadFailure = false; try { String keysetCursor = startCursor; @@ -979,6 +1103,7 @@ public class SearchIndexExecutor implements AutoCloseable { while (processed < recordLimit && !stopped.get()) { long backpressureWaitStart = System.currentTimeMillis(); + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 2000); while (isBackpressureActive()) { if (stopped.get()) { return; @@ -988,36 +1113,41 @@ public class SearchIndexExecutor implements AutoCloseable { LOG.warn("Backpressure wait timeout for {}, proceeding anyway", entityType); break; } - Thread.sleep(500); + Thread.sleep(backoff.nextDelay()); } try { - ResultList result = batchReader.readNextKeyset(keysetCursor); + ResultList result = readWithRetry(batchReader, keysetCursor, entityType); if (result == null || result.getData().isEmpty()) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (empty result)", + entityType, + processed, + recordLimit); break; } + if (!stopped.get()) { + IndexingTask task = new IndexingTask<>(entityType, result, processed); + taskQueue.put(task); + } + int readerSuccessCount = result.getData().size(); int readerFailedCount = listOrEmpty(result.getErrors()).size(); int readerWarningsCount = result.getWarningsCount() != null ? result.getWarningsCount() : 0; - updateReaderStats(readerSuccessCount, readerFailedCount, readerWarningsCount); - - StepStats batchStats = - new StepStats() - .withSuccessRecords(readerSuccessCount) - .withFailedRecords(readerFailedCount) - .withWarningRecords(readerWarningsCount); - updateStats(entityType, batchStats); - - if (!result.getData().isEmpty() && !stopped.get()) { - IndexingTask task = new IndexingTask<>(entityType, result, processed); - taskQueue.put(task); - } - processed += readerSuccessCount + readerFailedCount + readerWarningsCount; keysetCursor = result.getPaging() != null ? result.getPaging().getAfter() : null; if (keysetCursor == null) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (null cursor)", + entityType, + processed, + recordLimit); + break; + } + if (endCursor != null && hasReachedEndCursor(keysetCursor, endCursor)) { + LOG.debug("Reader for {} reached end cursor at processed={}", entityType, processed); break; } } catch (SearchIndexException e) { @@ -1047,7 +1177,7 @@ public class SearchIndexExecutor implements AutoCloseable { LOG.error("Error in keyset processing for {}", entityType, e); } } finally { - producerLatch.countDown(); + producerPhaser.arriveAndDeregister(); if (hadFailure) { AtomicInteger failures = entityBatchFailures.get(entityType); if (failures != null) { @@ -1068,8 +1198,8 @@ public class SearchIndexExecutor implements AutoCloseable { return; } - // Wait for backpressure to clear instead of dropping the batch long backpressureWaitStart = System.currentTimeMillis(); + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 2000); while (isBackpressureActive()) { if (stopped.get()) { return; @@ -1082,7 +1212,7 @@ public class SearchIndexExecutor implements AutoCloseable { currentOffset); break; } - Thread.sleep(500); + Thread.sleep(backoff.nextDelay()); } Source source = createSource(entityType); @@ -1142,7 +1272,60 @@ public class SearchIndexExecutor implements AutoCloseable { } } + private ResultList readWithRetry( + KeysetBatchReader batchReader, String keysetCursor, String entityType) + throws SearchIndexException, InterruptedException { + int maxRetryAttempts = 3; + long retryBackoffMs = 500; + for (int attempt = 0; attempt <= maxRetryAttempts; attempt++) { + try { + return batchReader.readNextKeyset(keysetCursor); + } catch (SearchIndexException e) { + if (attempt >= maxRetryAttempts || !isTransientReadError(e)) { + throw e; + } + long backoffDelay = retryBackoffMs * (1L << attempt); + LOG.warn( + "Transient read failure for {} (attempt {}/{}), retrying in {}ms", + entityType, + attempt + 1, + maxRetryAttempts, + backoffDelay); + Thread.sleep(Math.min(backoffDelay, 10_000)); + } + } + return null; + } + + private boolean isTransientReadError(SearchIndexException e) { + String msg = e.getMessage(); + if (msg == null) { + msg = ""; + } + String lower = msg.toLowerCase(); + return lower.contains("timeout") + || lower.contains("connection") + || lower.contains("pool exhausted") + || lower.contains("connectexception") + || lower.contains("sockettimeoutexception") + || lower.contains("remotetransportexception"); + } + private boolean isBackpressureActive() { + if (taskQueue != null) { + int size = taskQueue.size(); + int capacity = size + taskQueue.remainingCapacity(); + if (capacity > 0) { + int fillPercent = size * 100 / capacity; + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.updateQueueFillRatio(fillPercent); + } + if (fillPercent > 90) { + return true; + } + } + } if (lastBackpressureTime == 0) { return false; } @@ -1206,6 +1389,18 @@ public class SearchIndexExecutor implements AutoCloseable { return new PaginatedEntitiesSource( correctedEntityType, batchSize.get(), searchIndexFields, knownTotal); } else { + if (config != null) { + long startTs = config.getTimeSeriesStartTs(correctedEntityType); + if (startTs > 0) { + return new PaginatedEntityTimeSeriesSource( + correctedEntityType, + batchSize.get(), + searchIndexFields, + knownTotal, + startTs, + System.currentTimeMillis()); + } + } return new PaginatedEntityTimeSeriesSource( correctedEntityType, batchSize.get(), searchIndexFields, knownTotal); } @@ -1280,6 +1475,12 @@ public class SearchIndexExecutor implements AutoCloseable { sinkStats.setFailedRecords(0); jobDataStats.setSinkStats(sinkStats); + StepStats processStats = new StepStats(); + processStats.setTotalRecords(0); + processStats.setSuccessRecords(0); + processStats.setFailedRecords(0); + jobDataStats.setProcessStats(processStats); + return jobDataStats; } @@ -1303,6 +1504,13 @@ public class SearchIndexExecutor implements AutoCloseable { } else { repository = Entity.getEntityTimeSeriesRepository(entityType); } + if (config != null) { + long startTs = config.getTimeSeriesStartTs(correctedEntityType); + if (startTs > 0) { + long endTs = System.currentTimeMillis(); + return repository.getTimeSeriesDao().listCount(listFilter, startTs, endTs, false); + } + } return repository.getTimeSeriesDao().listCount(listFilter); } } catch (Exception e) { @@ -1311,18 +1519,6 @@ public class SearchIndexExecutor implements AutoCloseable { } } - private int getTotalLatchCount(Set entities, int fixedBatchSize) { - return entities.stream() - .mapToInt( - entityType -> { - int total = getTotalEntityRecords(entityType); - if (total <= 0) return 1; - return Math.min( - calculateNumberOfThreads(total, fixedBatchSize), MAX_READERS_PER_ENTITY); - }) - .sum(); - } - private int getTotalEntityRecords(String entityType) { if (stats.get() == null || stats.get().getEntityStats() == null @@ -1347,6 +1543,9 @@ public class SearchIndexExecutor implements AutoCloseable { } } + // Stats is published once via stats.set(initializeTotalRecords(...)) and all subsequent + // mutations operate on that same mutable object under synchronized methods. + synchronized void updateStats(String entityType, StepStats currentEntityStats) { Stats jobDataStats = stats.get(); if (jobDataStats == null) { @@ -1355,7 +1554,6 @@ public class SearchIndexExecutor implements AutoCloseable { updateEntityStats(jobDataStats, entityType, currentEntityStats); updateJobStats(jobDataStats); - stats.set(jobDataStats); } synchronized void updateReaderStats(int successCount, int failedCount, int warningsCount) { @@ -1379,8 +1577,6 @@ public class SearchIndexExecutor implements AutoCloseable { readerStats.setSuccessRecords(currentSuccess + successCount); readerStats.setFailedRecords(currentFailed + failedCount); readerStats.setWarningRecords(currentWarnings + warningsCount); - - stats.set(jobDataStats); } synchronized void updateSinkTotalSubmitted(int submittedCount) { @@ -1398,8 +1594,6 @@ public class SearchIndexExecutor implements AutoCloseable { int currentTotal = sinkStats.getTotalRecords() != null ? sinkStats.getTotalRecords() : 0; sinkStats.setTotalRecords(currentTotal + submittedCount); - - stats.set(jobDataStats); } synchronized void syncSinkStatsFromBulkSink() { @@ -1423,6 +1617,8 @@ public class SearchIndexExecutor implements AutoCloseable { jobDataStats.setSinkStats(sinkStats); } + sinkStats.setTotalRecords( + bulkSinkStats.getTotalRecords() != null ? bulkSinkStats.getTotalRecords() : 0); sinkStats.setSuccessRecords( bulkSinkStats.getSuccessRecords() != null ? bulkSinkStats.getSuccessRecords() : 0); sinkStats.setFailedRecords( @@ -1435,7 +1631,11 @@ public class SearchIndexExecutor implements AutoCloseable { jobDataStats.setVectorStats(vectorStats); } - stats.set(jobDataStats); + // Sync process stats if available + StepStats processStats = searchIndexSink.getProcessStats(); + if (processStats != null) { + jobDataStats.setProcessStats(processStats); + } } private void updateEntityStats(Stats statsObj, String entityType, StepStats currentEntityStats) { @@ -1505,16 +1705,20 @@ public class SearchIndexExecutor implements AutoCloseable { private void closeSinkIfNeeded() throws IOException { if (searchIndexSink != null && sinkClosed.compareAndSet(false, true)) { - // Check for pending vector tasks before closing int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); if (pendingVectorTasks > 0) { LOG.info( "Waiting for {} pending vector embedding tasks to complete before closing", pendingVectorTasks); + VectorCompletionResult vcResult = searchIndexSink.awaitVectorCompletionWithDetails(300); + LOG.info( + "Vector completion: completed={}, pending={}, waited={}ms", + vcResult.completed(), + vcResult.pendingTaskCount(), + vcResult.waitedMillis()); } LOG.info("Forcing final flush of bulk processor and vector embeddings"); - // close() internally calls awaitVectorCompletion() first, then flushes search index searchIndexSink.close(); syncSinkStatsFromBulkSink(); } @@ -1530,7 +1734,6 @@ public class SearchIndexExecutor implements AutoCloseable { Stats currentStats = stats.get(); if (currentStats != null) { StatsReconciler.reconcile(currentStats); - stats.set(currentStats); } long endTime = System.currentTimeMillis(); @@ -1580,17 +1783,40 @@ public class SearchIndexExecutor implements AutoCloseable { listeners.onJobStopped(stats.get()); - if (taskQueue != null) { - taskQueue.clear(); - for (int i = 0; i < 10; i++) { - taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); - } + if (searchIndexSink != null) { + LOG.info( + "Stopping executor: flushing sink ({} active bulk requests)", + searchIndexSink.getActiveBulkRequestCount()); + searchIndexSink.flushAndAwait(10); + } + + int dropped = taskQueue != null ? taskQueue.size() : 0; + if (dropped > 0) { + LOG.warn("Dropping {} queued tasks during shutdown", dropped); } shutdownExecutor(producerExecutor, "producer"); - shutdownExecutor(consumerExecutor, "consumer"); shutdownExecutor(jobExecutor, "job"); + if (taskQueue != null) { + taskQueue.clear(); + for (int i = 0; i < MAX_CONSUMER_THREADS; i++) { + taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + if (consumerExecutor != null && !consumerExecutor.isShutdown()) { + consumerExecutor.shutdown(); + try { + if (!consumerExecutor.awaitTermination(5, TimeUnit.SECONDS)) { + consumerExecutor.shutdownNow(); + LOG.warn("Consumer executor did not terminate within 5s, forced shutdown"); + } + } catch (InterruptedException e) { + consumerExecutor.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + LOG.info("Reindexing executor stopped"); } @@ -1700,6 +1926,7 @@ public class SearchIndexExecutor implements AutoCloseable { if (statsManager != null) { statsManager.flushAll(); } + sinkTrackers.values().forEach(StageStatsTracker::flush); stop(); cleanup(); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java new file mode 100644 index 00000000000..d347514bdb6 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java @@ -0,0 +1,41 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Optional; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; + +public class SingleServerIndexingStrategy implements IndexingStrategy { + + private final SearchIndexExecutor executor; + + public SingleServerIndexingStrategy( + CollectionDAO collectionDAO, SearchRepository searchRepository) { + this.executor = new SearchIndexExecutor(collectionDAO, searchRepository); + } + + @Override + public void addListener(ReindexingProgressListener listener) { + executor.addListener(listener); + } + + @Override + public ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context) { + return executor.execute(config, context); + } + + @Override + public Optional getStats() { + return Optional.ofNullable(executor.getStats().get()); + } + + @Override + public void stop() { + executor.stop(); + } + + @Override + public boolean isStopped() { + return executor.isStopped(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/StatsReconciler.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/StatsReconciler.java index e05e0438ef4..24c9dbf5ae5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/StatsReconciler.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/StatsReconciler.java @@ -30,11 +30,26 @@ public class StatsReconciler { int sinkFailed = safeGet(sinkStats.getFailedRecords()); int sinkWarnings = safeGet(sinkStats.getWarningRecords()); + // Reconcile entity-level totals + if (stats.getEntityStats() != null + && stats.getEntityStats().getAdditionalProperties() != null) { + int reconciledTotal = 0; + for (StepStats es : stats.getEntityStats().getAdditionalProperties().values()) { + int actual = safeGet(es.getSuccessRecords()) + safeGet(es.getFailedRecords()); + if (actual > safeGet(es.getTotalRecords())) { + es.setTotalRecords(actual); + } + reconciledTotal += safeGet(es.getTotalRecords()); + } + if (reconciledTotal > readerTotal) { + readerStats.setTotalRecords(reconciledTotal); + readerTotal = reconciledTotal; + } + } + int jobSuccess = sinkSuccess; int jobFailed = readerFailed + sinkFailed; int jobTotal = readerTotal; - // Warnings are informational - use reader warnings as the primary source - // (entities with stale references that were still indexed) int jobWarnings = readerWarnings; jobStats.setTotalRecords(jobTotal); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/VectorCompletionResult.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/VectorCompletionResult.java new file mode 100644 index 00000000000..5be32a0bca4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/VectorCompletionResult.java @@ -0,0 +1,12 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +public record VectorCompletionResult(boolean completed, int pendingTaskCount, long waitedMillis) { + + public static VectorCompletionResult success(long waitedMillis) { + return new VectorCompletionResult(true, 0, waitedMillis); + } + + public static VectorCompletionResult timeout(int pendingCount, long waitedMillis) { + return new VectorCompletionResult(false, pendingCount, waitedMillis); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java index a6de639dd66..6bf2f7d31aa 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java @@ -269,9 +269,13 @@ public class DistributedJobParticipant implements Managed { // Set up failure callback on bulk sink to record sink failures final IndexingFailureRecorder recorder = failureRecorder; bulkSink.setFailureCallback( - (entityType, entityId, entityFqn, errorMessage) -> { + (entityType, entityId, entityFqn, errorMessage, stage) -> { if (recorder != null) { - recorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + if (stage == IndexingFailureRecorder.FailureStage.PROCESS) { + recorder.recordProcessFailure(entityType, entityId, entityFqn, errorMessage); + } else { + recorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + } } }); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java index f0a2514b662..c621f7aad9f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java @@ -20,10 +20,12 @@ import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.system.EventPublisherJob; import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.jdbi3.CollectionDAO.SearchIndexJobDAO; import org.openmetadata.service.jdbi3.CollectionDAO.SearchIndexJobDAO.SearchIndexJobRecord; @@ -85,6 +87,9 @@ public class DistributedSearchIndexCoordinator { private final String serverId; private EntityCompletionTracker entityTracker; + /** Monotonic counter to guarantee unique claimedAt values across concurrent worker threads. */ + private final AtomicLong claimCounter = new AtomicLong(0); + public DistributedSearchIndexCoordinator(CollectionDAO collectionDAO) { this.collectionDAO = collectionDAO; this.partitionCalculator = new PartitionCalculator(); @@ -121,12 +126,20 @@ public class DistributedSearchIndexCoordinator { */ public SearchIndexJob createJob( Set entities, EventPublisherJob jobConfiguration, String createdBy) { + return createJob(entities, jobConfiguration, createdBy, null); + } + + public SearchIndexJob createJob( + Set entities, + EventPublisherJob jobConfiguration, + String createdBy, + ReindexingConfiguration reindexConfig) { UUID jobId = UUID.randomUUID(); long now = System.currentTimeMillis(); - // Calculate entity statistics - Map entityCounts = partitionCalculator.getEntityCounts(entities); + // Calculate entity statistics (with time-series date filtering if config is provided) + Map entityCounts = partitionCalculator.getEntityCounts(entities, reindexConfig); long totalRecords = entityCounts.values().stream().mapToLong(Long::longValue).sum(); // Build entity stats map @@ -183,6 +196,10 @@ public class DistributedSearchIndexCoordinator { * @return Updated job with partition information */ public SearchIndexJob initializePartitions(UUID jobId) { + return initializePartitions(jobId, null); + } + + public SearchIndexJob initializePartitions(UUID jobId, ReindexingConfiguration reindexConfig) { SearchIndexJobDAO jobDAO = collectionDAO.searchIndexJobDAO(); SearchIndexPartitionDAO partitionDAO = collectionDAO.searchIndexPartitionDAO(); @@ -196,9 +213,9 @@ public class DistributedSearchIndexCoordinator { // Get entity types from job configuration Set entityTypes = Set.copyOf(job.getJobConfiguration().getEntities()); - // Calculate partitions + // Calculate partitions (with date filtering for time series if config provided) List partitions = - partitionCalculator.calculatePartitions(jobId, entityTypes); + partitionCalculator.calculatePartitions(jobId, entityTypes, reindexConfig); if (partitions.isEmpty()) { LOG.warn( @@ -270,10 +287,22 @@ public class DistributedSearchIndexCoordinator { } } + // Reconcile totalRecords from actual partitions (accounts for time-series filtering) + long actualTotalRecords = + partitions.stream().mapToLong(SearchIndexPartition::getEstimatedCount).sum(); + if (actualTotalRecords != job.getTotalRecords()) { + LOG.info( + "Reconciled totalRecords for job {}: {} → {} (after partition calculation)", + jobId, + job.getTotalRecords(), + actualTotalRecords); + } + // Update job status SearchIndexJob updatedJob = job.toBuilder() .status(IndexJobStatus.READY) + .totalRecords(actualTotalRecords) .entityStats(updatedStats) .updatedAt(System.currentTimeMillis()) .build(); @@ -313,7 +342,9 @@ public class DistributedSearchIndexCoordinator { return Optional.empty(); } - long claimTime = System.currentTimeMillis(); + // Ensure unique claimTime per call so concurrent claims on the same server are distinguishable. + // The counter suffix keeps values within normal epoch-millis range while preventing collisions. + long claimTime = uniqueClaimTime(); // Atomically claim a partition - FOR UPDATE SKIP LOCKED ensures no race condition int claimed = partitionDAO.claimNextPartitionAtomic(jobId.toString(), serverId, claimTime); @@ -322,9 +353,9 @@ public class DistributedSearchIndexCoordinator { return Optional.empty(); } - // Fetch the partition we just claimed + // Fetch the partition we just claimed using the unique claimTime SearchIndexPartitionRecord record = - partitionDAO.findLatestClaimedPartition(jobId.toString(), serverId); + partitionDAO.findLatestClaimedPartition(jobId.toString(), serverId, claimTime); if (record == null) { LOG.warn("Claimed partition but couldn't find it - this shouldn't happen"); return Optional.empty(); @@ -343,6 +374,18 @@ public class DistributedSearchIndexCoordinator { return Optional.of(partition); } + /** + * Generates a unique claimedAt timestamp that stays close to real wall-clock time but never + * repeats, even when called concurrently from multiple worker threads. The counter suffix is + * added in the sub-millisecond range so stale-detection logic (which compares against + * System.currentTimeMillis()) continues to work correctly. + */ + private long uniqueClaimTime() { + long millis = System.currentTimeMillis(); + long seq = claimCounter.incrementAndGet() % 1000; + return millis + seq; + } + /** * Update partition progress. * diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java index 1af8ecd1aef..36dc3f2b6fd 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java @@ -139,7 +139,7 @@ public class DistributedSearchIndexExecutor { public DistributedSearchIndexExecutor(CollectionDAO collectionDAO, int partitionSize) { this.collectionDAO = collectionDAO; - PartitionCalculator calculator = new PartitionCalculator(partitionSize); + PartitionCalculator calculator = new PartitionCalculator(partitionSize, MAX_WORKER_THREADS); this.coordinator = new DistributedSearchIndexCoordinator(collectionDAO, calculator); this.recoveryManager = new JobRecoveryManager(collectionDAO, partitionSize); this.serverId = ServerIdentityResolver.getInstance().getServerId(); @@ -207,7 +207,10 @@ public class DistributedSearchIndexExecutor { * @return The created job */ public SearchIndexJob createJob( - Set entities, EventPublisherJob jobConfiguration, String createdBy) { + Set entities, + EventPublisherJob jobConfiguration, + String createdBy, + ReindexingConfiguration reindexConfig) { LOG.info("Creating distributed indexing job for {} entity types", entities.size()); @@ -240,11 +243,12 @@ public class DistributedSearchIndexExecutor { } try { - // Create the job - SearchIndexJob job = coordinator.createJob(entities, jobConfiguration, createdBy); + // Create the job (pass reindexConfig so time-series date filtering is applied to totals) + SearchIndexJob job = + coordinator.createJob(entities, jobConfiguration, createdBy, reindexConfig); - // Initialize partitions - currentJob = coordinator.initializePartitions(job.getId()); + // Initialize partitions (with date filtering for time series entities) + currentJob = coordinator.initializePartitions(job.getId(), reindexConfig); // Atomically transfer lock to real job ID boolean transferred = coordinator.transferReindexLock(tempJobId, currentJob.getId()); @@ -306,7 +310,10 @@ public class DistributedSearchIndexExecutor { * @return Execution result with statistics */ public ExecutionResult execute( - BulkSink bulkSink, ReindexContext recreateContext, boolean recreateIndex) { + BulkSink bulkSink, + ReindexContext recreateContext, + boolean recreateIndex, + ReindexingConfiguration reindexConfig) { if (currentJob == null) { throw new IllegalStateException("No job to execute - call createJob() or joinJob() first"); @@ -342,11 +349,9 @@ public class DistributedSearchIndexExecutor { // Notify listeners that job has started listeners.onJobStarted(jobContext); - // Notify listeners with configuration - if (currentJob.getJobConfiguration() != null) { - ReindexingConfiguration config = - ReindexingConfiguration.from(currentJob.getJobConfiguration()); - listeners.onJobConfigured(jobContext, config); + // Notify listeners with auto-tuned configuration + if (reindexConfig != null) { + listeners.onJobConfigured(jobContext, reindexConfig); } // Create stats aggregator with app context for proper WebSocket matching @@ -373,9 +378,13 @@ public class DistributedSearchIndexExecutor { // Set up failure callback on the sink to record sink failures bulkSink.setFailureCallback( - (entityType, entityId, entityFqn, errorMessage) -> { + (entityType, entityId, entityFqn, errorMessage, stage) -> { if (failureRecorder != null) { - failureRecorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + if (stage == IndexingFailureRecorder.FailureStage.PROCESS) { + failureRecorder.recordProcessFailure(entityType, entityId, entityFqn, errorMessage); + } else { + failureRecorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + } } }); @@ -402,13 +411,13 @@ public class DistributedSearchIndexExecutor { .name("partition-heartbeat-" + jobId.toString().substring(0, 8)) .start(() -> runPartitionHeartbeatLoop()); - // Calculate worker threads based on configuration - int numWorkers = - Math.min( - currentJob.getJobConfiguration().getConsumerThreads() != null - ? currentJob.getJobConfiguration().getConsumerThreads() - : 4, - MAX_WORKER_THREADS); + // Calculate worker threads from auto-tuned configuration + int numWorkers = Math.min(Math.max(1, reindexConfig.consumerThreads()), MAX_WORKER_THREADS); + LOG.info( + "Distributed executor using {} workers, batch size {} (autoTune={})", + numWorkers, + reindexConfig.batchSize(), + reindexConfig.autoTune()); workerExecutor = Executors.newFixedThreadPool( @@ -419,10 +428,7 @@ public class DistributedSearchIndexExecutor { CountDownLatch workerLatch = new CountDownLatch(numWorkers); // Start worker threads that continuously claim and process partitions - int batchSize = - currentJob.getJobConfiguration().getBatchSize() != null - ? currentJob.getJobConfiguration().getBatchSize() - : 500; + int batchSize = reindexConfig.batchSize(); for (int i = 0; i < numWorkers; i++) { final int workerId = i; @@ -436,7 +442,8 @@ public class DistributedSearchIndexExecutor { recreateContext, recreateIndex, totalSuccess, - totalFailed); + totalFailed, + reindexConfig); } finally { workerLatch.countDown(); } @@ -458,6 +465,18 @@ public class DistributedSearchIndexExecutor { // entity types have 0 records), so no partition completion ever triggers the check. coordinator.checkAndUpdateJobCompletion(jobId); + // Final reconciliation pass: catch ALL participant-server completions before + // the stale-reclaimer is killed. Participant workers may have finished partitions + // that were never reconciled by the stale-reclaimer's periodic loop. + if (entityTracker != null && recreateContext != null) { + LOG.info("Running final DB reconciliation for job {}", jobId); + List allPartitions = coordinator.getPartitions(jobId, null); + entityTracker.reconcileFromDatabase(allPartitions); + LOG.info( + "Final reconciliation complete - promoted entities: {}", + entityTracker.getPromotedEntities()); + } + } catch (InterruptedException e) { Thread.currentThread().interrupt(); LOG.warn("Execution interrupted for job {}", jobId); @@ -538,13 +557,20 @@ public class DistributedSearchIndexExecutor { ReindexContext recreateContext, boolean recreateIndex, AtomicLong totalSuccess, - AtomicLong totalFailed) { + AtomicLong totalFailed, + ReindexingConfiguration reindexConfig) { LOG.info("Worker {} starting for job {}", workerId, currentJob.getId()); PartitionWorker worker = new PartitionWorker( - coordinator, bulkSink, batchSize, recreateContext, recreateIndex, failureRecorder); + coordinator, + bulkSink, + batchSize, + recreateContext, + recreateIndex, + failureRecorder, + reindexConfig); synchronized (activeWorkers) { activeWorkers.add(worker); @@ -991,12 +1017,22 @@ public class DistributedSearchIndexExecutor { } try { + String canonicalIndex = recreateContext.getCanonicalIndex(entityType).orElse(null); + String originalIndex = recreateContext.getOriginalIndex(entityType).orElse(null); + + LOG.debug( + "Promoting entity '{}': success={}, canonicalIndex={}, stagedIndex={}", + entityType, + success, + canonicalIndex, + stagedIndexOpt.get()); + EntityReindexContext entityContext = EntityReindexContext.builder() .entityType(entityType) - .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) - .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) - .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .originalIndex(originalIndex) + .canonicalIndex(canonicalIndex) + .activeIndex(originalIndex) .stagedIndex(stagedIndexOpt.get()) .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) .existingAliases(recreateContext.getExistingAliases(entityType)) diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/EntityCompletionTracker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/EntityCompletionTracker.java index f52e23fcf30..3149c029860 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/EntityCompletionTracker.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/EntityCompletionTracker.java @@ -155,18 +155,33 @@ public class EntityCompletionTracker { continue; } - boolean allDone = + long completedCount = entityPartitions.stream() - .allMatch( + .filter( p -> p.getStatus() == PartitionStatus.COMPLETED - || p.getStatus() == PartitionStatus.FAILED); + || p.getStatus() == PartitionStatus.FAILED) + .count(); + boolean allDone = completedCount == entityPartitions.size(); + + if (!allDone) { + Map statusCounts = + entityPartitions.stream() + .collect( + Collectors.groupingBy(SearchIndexPartition::getStatus, Collectors.counting())); + LOG.debug( + "Reconcile: entity '{}' not all done: {}/{} complete, statuses={}", + entityType, + completedCount, + entityPartitions.size(), + statusCounts); + } if (allDone && !entityPartitions.isEmpty()) { boolean hasFailed = entityPartitions.stream().anyMatch(p -> p.getStatus() == PartitionStatus.FAILED); - LOG.info( + LOG.debug( "DB reconciliation: entity '{}' all {} partitions done (hasFailed={}, job {})", entityType, entityPartitions.size(), @@ -182,16 +197,27 @@ public class EntityCompletionTracker { if (promotedEntities.add(entityType)) { boolean success = !hasFailed; - LOG.info( - "Entity '{}' all partitions complete (success={}, job {})", entityType, success, jobId); + LOG.debug( + "Entity '{}' all partitions complete (success={}, hasFailed={}, job {})", + entityType, + success, + hasFailed, + jobId); if (onEntityComplete != null) { try { onEntityComplete.accept(entityType, success); } catch (Exception e) { - LOG.error("Error in entity completion callback for '{}' (job {})", entityType, jobId, e); + LOG.error( + "Error in entity completion callback for '{}' (job {}). " + + "Entity is STILL in promotedEntities - will be SKIPPED by finalization!", + entityType, + jobId, + e); } } + } else { + LOG.debug("Entity '{}' already in promotedEntities, skipping (job {})", entityType, jobId); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java index ad1ebc41592..225424d7305 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java @@ -22,6 +22,8 @@ import java.util.UUID; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.type.Include; import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.EntityPriority; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; import org.openmetadata.service.jdbi3.EntityRepository; import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; import org.openmetadata.service.jdbi3.ListFilter; @@ -79,38 +81,6 @@ public class PartitionCalculator { Map.entry("queryCostRecord", 0.3) // Time series, simple structure ); - /** - * Priority ordering for entity types during indexing. Higher priority entities should be indexed - * first as they may be referenced by others. This ensures that when indexing tables, their parent - * databases and schemas already exist in the search index. - */ - private static final Map ENTITY_PRIORITY = - Map.ofEntries( - Map.entry("databaseService", 100), - Map.entry("messagingService", 100), - Map.entry("dashboardService", 100), - Map.entry("pipelineService", 100), - Map.entry("mlmodelService", 100), - Map.entry("storageService", 100), - Map.entry("database", 90), - Map.entry("databaseSchema", 80), - Map.entry("glossary", 70), - Map.entry("classification", 70), - Map.entry("team", 65), - Map.entry("user", 60), - Map.entry("table", 50), - Map.entry("dashboard", 50), - Map.entry("pipeline", 50), - Map.entry("mlmodel", 50), - Map.entry("topic", 50), - Map.entry("container", 50), - Map.entry("glossaryTerm", 45), - Map.entry("tag", 40), - Map.entry("testCase", 30), - Map.entry("testCaseResult", 20), - Map.entry("testCaseResolutionStatus", 20), - Map.entry("queryCostRecord", 10)); - /** Time series entity types */ private static final Set TIME_SERIES_ENTITIES = Set.of( @@ -124,13 +94,19 @@ public class PartitionCalculator { "aggregatedCostAnalysisReportData"); private final int partitionSize; + private final int minPartitionsPerEntity; public PartitionCalculator() { - this(DEFAULT_PARTITION_SIZE); + this(DEFAULT_PARTITION_SIZE, 1); } public PartitionCalculator(int partitionSize) { + this(partitionSize, 1); + } + + public PartitionCalculator(int partitionSize, int minPartitionsPerEntity) { this.partitionSize = Math.clamp(partitionSize, MIN_PARTITION_SIZE, MAX_PARTITION_SIZE); + this.minPartitionsPerEntity = Math.max(1, minPartitionsPerEntity); } /** @@ -144,10 +120,16 @@ public class PartitionCalculator { * @throws IllegalStateException if partition count would exceed safe limits */ public List calculatePartitions(UUID jobId, Set entityTypes) { + return calculatePartitions(jobId, entityTypes, null); + } + + public List calculatePartitions( + UUID jobId, Set entityTypes, ReindexingConfiguration reindexConfig) { List partitions = new ArrayList<>(); for (String entityType : entityTypes) { - List entityPartitions = calculatePartitionsForEntity(jobId, entityType); + List entityPartitions = + calculatePartitionsForEntity(jobId, entityType, reindexConfig); partitions.addAll(entityPartitions); if (partitions.size() > MAX_TOTAL_PARTITIONS) { @@ -174,14 +156,19 @@ public class PartitionCalculator { * @return List of partitions for this entity type */ public List calculatePartitionsForEntity(UUID jobId, String entityType) { - long totalCount = getEntityCount(entityType); + return calculatePartitionsForEntity(jobId, entityType, null); + } + + public List calculatePartitionsForEntity( + UUID jobId, String entityType, ReindexingConfiguration reindexConfig) { + long totalCount = getEntityCount(entityType, reindexConfig); if (totalCount == 0) { LOG.debug("No entities found for type: {}", entityType); return List.of(); } double complexityFactor = ENTITY_COMPLEXITY_FACTORS.getOrDefault(entityType, 1.0); - int priority = ENTITY_PRIORITY.getOrDefault(entityType, 50); + int priority = EntityPriority.getNumericPriority(entityType); // Adjust partition size based on complexity - more complex entities get smaller partitions long adjustedPartitionSizeLong = (long) (partitionSize / complexityFactor); @@ -191,6 +178,13 @@ public class PartitionCalculator { long numPartitionsLong = (totalCount + adjustedPartitionSizeLong - 1) / adjustedPartitionSizeLong; + // Ensure minimum partitions so all workers stay busy (e.g. testCaseResult with + // only 4 partitions leaves 6 of 10 workers idle for minutes) + if (numPartitionsLong < minPartitionsPerEntity && totalCount >= minPartitionsPerEntity) { + numPartitionsLong = minPartitionsPerEntity; + adjustedPartitionSizeLong = (totalCount + numPartitionsLong - 1) / numPartitionsLong; + } + // Enforce per-entity-type limit and adjust partition size if needed if (numPartitionsLong > MAX_PARTITIONS_PER_ENTITY_TYPE) { LOG.warn( @@ -253,10 +247,14 @@ public class PartitionCalculator { * @return Total count of entities */ public long getEntityCount(String entityType) { + return getEntityCount(entityType, null); + } + + public long getEntityCount(String entityType, ReindexingConfiguration reindexConfig) { try { long count; if (TIME_SERIES_ENTITIES.contains(entityType)) { - count = getTimeSeriesEntityCount(entityType); + count = getTimeSeriesEntityCount(entityType, reindexConfig); } else { count = getRegularEntityCount(entityType); } @@ -270,10 +268,10 @@ public class PartitionCalculator { private long getRegularEntityCount(String entityType) { EntityRepository repository = Entity.getEntityRepository(entityType); - return repository.getDao().listTotalCount(); + return repository.getDao().listCount(new ListFilter(Include.ALL)); } - private long getTimeSeriesEntityCount(String entityType) { + private long getTimeSeriesEntityCount(String entityType, ReindexingConfiguration reindexConfig) { ListFilter listFilter = new ListFilter(Include.ALL); EntityTimeSeriesRepository repository; @@ -284,6 +282,21 @@ public class PartitionCalculator { repository = Entity.getEntityTimeSeriesRepository(entityType); } + if (reindexConfig != null) { + long startTs = reindexConfig.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + long endTs = System.currentTimeMillis(); + long count = repository.getTimeSeriesDao().listCount(listFilter, startTs, endTs, false); + LOG.info( + "Time series date filter for {}: last {} days → {} records (was {} total)", + entityType, + reindexConfig.timeSeriesMaxDays(), + count, + repository.getTimeSeriesDao().listCount(listFilter)); + return count; + } + } + return repository.getTimeSeriesDao().listCount(listFilter); } @@ -298,9 +311,14 @@ public class PartitionCalculator { * @return Map of entity type to count */ public Map getEntityCounts(Set entityTypes) { + return getEntityCounts(entityTypes, null); + } + + public Map getEntityCounts( + Set entityTypes, ReindexingConfiguration reindexConfig) { Map counts = new HashMap<>(); for (String entityType : entityTypes) { - counts.put(entityType, getEntityCount(entityType)); + counts.put(entityType, getEntityCount(entityType, reindexConfig)); } return counts; } @@ -312,7 +330,7 @@ public class PartitionCalculator { * @return Priority value (higher = processed first) */ public int getEntityPriority(String entityType) { - return ENTITY_PRIORITY.getOrDefault(entityType, 50); + return EntityPriority.getNumericPriority(entityType); } /** diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java index c37e25662c3..5ddb8dc25eb 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java @@ -28,11 +28,13 @@ import org.apache.commons.lang3.exception.ExceptionUtils; import org.openmetadata.schema.EntityInterface; import org.openmetadata.schema.EntityTimeSeriesInterface; import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.system.EntityError; import org.openmetadata.schema.type.Include; import org.openmetadata.schema.utils.ResultList; import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; import org.openmetadata.service.exception.SearchIndexException; import org.openmetadata.service.jdbi3.ListFilter; @@ -77,6 +79,12 @@ public class PartitionWorker { /** Progress update interval (every N entities) */ private static final int PROGRESS_UPDATE_INTERVAL = 100; + /** Overall deadline for waiting on sink operations to complete */ + private static final long SINK_WAIT_DEADLINE_MS = 300_000; + + /** Timeout per flush cycle when retrying sink completion */ + private static final int FLUSH_CYCLE_SECONDS = 30; + private final DistributedSearchIndexCoordinator coordinator; private final BulkSink searchIndexSink; private final int batchSize; @@ -84,6 +92,7 @@ public class PartitionWorker { private final boolean recreateIndex; private final AtomicBoolean stopped = new AtomicBoolean(false); private final IndexingFailureRecorder failureRecorder; + private final ReindexingConfiguration reindexConfig; public PartitionWorker( DistributedSearchIndexCoordinator coordinator, @@ -91,7 +100,7 @@ public class PartitionWorker { int batchSize, ReindexContext recreateContext, boolean recreateIndex) { - this(coordinator, searchIndexSink, batchSize, recreateContext, recreateIndex, null); + this(coordinator, searchIndexSink, batchSize, recreateContext, recreateIndex, null, null); } public PartitionWorker( @@ -101,12 +110,31 @@ public class PartitionWorker { ReindexContext recreateContext, boolean recreateIndex, IndexingFailureRecorder failureRecorder) { + this( + coordinator, + searchIndexSink, + batchSize, + recreateContext, + recreateIndex, + failureRecorder, + null); + } + + public PartitionWorker( + DistributedSearchIndexCoordinator coordinator, + BulkSink searchIndexSink, + int batchSize, + ReindexContext recreateContext, + boolean recreateIndex, + IndexingFailureRecorder failureRecorder, + ReindexingConfiguration reindexConfig) { this.coordinator = coordinator; this.searchIndexSink = searchIndexSink; this.batchSize = batchSize; this.recreateContext = recreateContext; this.recreateIndex = recreateIndex; this.failureRecorder = failureRecorder; + this.reindexConfig = reindexConfig; } /** @@ -154,19 +182,25 @@ public class PartitionWorker { // Initialize keyset cursor for efficient pagination (avoids OFFSET degradation) long cursorInitStart = System.currentTimeMillis(); String keysetCursor = initializeKeysetCursor(entityType, rangeStart); - LOG.info( - "[PERF] initializeKeysetCursor for {} offset={} took {}ms", + LOG.debug( + "initializeKeysetCursor for {} offset={} took {}ms", entityType, rangeStart, System.currentTimeMillis() - cursorInitStart); // Process in batches - while (currentOffset < rangeEnd && !stopped.get()) { + while (currentOffset < rangeEnd + && !stopped.get() + && !Thread.currentThread().isInterrupted()) { int currentBatchSize = (int) Math.min(batchSize, rangeEnd - currentOffset); try { BatchResult batchResult = processBatch(entityType, keysetCursor, currentBatchSize, statsTracker); + // Check for stop/interrupt after DB read completes + if (stopped.get() || Thread.currentThread().isInterrupted()) { + break; + } successCount.addAndGet(batchResult.successCount()); failedCount.addAndGet(batchResult.failedCount()); warningsCount.addAndGet(batchResult.warningsCount()); @@ -191,9 +225,14 @@ public class PartitionWorker { keysetCursor = initializeKeysetCursor(entityType, currentOffset); if (keysetCursor == null) { LOG.debug( - "No more data at offset {} (rangeEnd: {}), finishing partition early", + "{} partition {} data exhausted at offset {} (rangeEnd: {}), " + + "missing {} records. processedCount={}", + entityType, + partition.getId(), currentOffset, - rangeEnd); + rangeEnd, + rangeEnd - currentOffset, + processedCount.get()); break; } } @@ -272,11 +311,13 @@ public class PartitionWorker { // the coordinator may aggregate stats before they're written to the database long waitStart = System.currentTimeMillis(); waitForSinkOperations(statsTracker); - LOG.info("[PERF] waitForSinkOperations took {}ms", System.currentTimeMillis() - waitStart); + LOG.debug("waitForSinkOperations took {}ms", System.currentTimeMillis() - waitStart); // Mark partition as completed (stats are now in the database) coordinator.completePartition(partition.getId(), successCount.get(), failedCount.get()); + long expectedRecords = rangeEnd - rangeStart; + long actualProcessed = successCount.get() + failedCount.get(); LOG.info( "Completed partition {} for entity type {} (success: {}, failed: {}, readerFailed: {}, warnings: {})", partition.getId(), @@ -285,6 +326,18 @@ public class PartitionWorker { failedCount.get(), readerFailedCount.get(), warningsCount.get()); + if (actualProcessed < expectedRecords) { + LOG.debug( + "{} partition {} processed fewer records than expected: " + + "actual={}, expected={}, gap={}, range=[{},{})", + entityType, + partition.getId(), + actualProcessed, + expectedRecords, + expectedRecords - actualProcessed, + rangeStart, + rangeEnd); + } return new PartitionResult( successCount.get(), @@ -322,7 +375,7 @@ public class PartitionWorker { private void waitForSinkOperations(StageStatsTracker statsTracker) { // Flush the bulk processor to send any pending documents immediately // Without this, documents wait for the periodic flush interval (5 seconds) - searchIndexSink.flushAndAwait(30); + searchIndexSink.flushAndAwait(FLUSH_CYCLE_SECONDS); // Check if there are pending vector tasks - if so, we need a longer timeout int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); @@ -334,7 +387,6 @@ public class PartitionWorker { pendingVectorTasks, statsTracker.getEntityType()); - // Wait for vector operations to complete first (up to 120 seconds for vectors) boolean vectorComplete = searchIndexSink.awaitVectorCompletion(120); if (!vectorComplete) { LOG.warn( @@ -344,15 +396,59 @@ public class PartitionWorker { } } - // Now wait for the stats tracker to have all callbacks accounted for - // Use a longer timeout if we had vector tasks since callbacks may be delayed - long statsTimeout = hasVectorTasks ? 60000 : 30000; - boolean statsComplete = statsTracker.awaitSinkCompletion(statsTimeout); - if (!statsComplete) { + // Wait for all sink callbacks with retries. The bulk processor is shared across + // partition workers, so slow batches from other entity types (e.g. testCaseResult + // writes taking 70+ seconds) can delay our callbacks. Instead of a single fixed + // timeout, retry flush cycles until all pending operations complete. + long deadline = System.currentTimeMillis() + SINK_WAIT_DEADLINE_MS; + int retryCount = 0; + long previousPending = statsTracker.getPendingSinkOps(); + int staleRetries = 0; + + while (statsTracker.getPendingSinkOps() > 0 && System.currentTimeMillis() < deadline) { + long remainingMs = deadline - System.currentTimeMillis(); + long waitMs = Math.min(30_000, remainingMs); + + if (statsTracker.awaitSinkCompletion(waitMs)) { + break; + } + + if (statsTracker.getPendingSinkOps() > 0 && System.currentTimeMillis() < deadline) { + retryCount++; + long currentPending = statsTracker.getPendingSinkOps(); + LOG.info( + "Retry {} - {} sink operations still pending for entity {}, re-flushing bulk processor", + retryCount, + currentPending, + statsTracker.getEntityType()); + searchIndexSink.flushAndAwait(FLUSH_CYCLE_SECONDS); + + if (currentPending == previousPending) { + staleRetries++; + if (staleRetries >= 3) { + LOG.warn( + "Pending sink ops stuck at {} for entity {} after {} retries with no progress. " + + "Reconciling early (callbacks likely lost).", + currentPending, + statsTracker.getEntityType(), + staleRetries); + break; + } + } else { + staleRetries = 0; + } + previousPending = currentPending; + } + } + + if (statsTracker.getPendingSinkOps() > 0) { LOG.warn( - "Timed out waiting for sink stats completion, {} operations still pending for entity {}", + "Reconciling {} pending sink operations after {} retries for entity {} " + + "(bulk processor was flushed, treating as successful)", statsTracker.getPendingSinkOps(), + retryCount, statsTracker.getEntityType()); + statsTracker.reconcilePendingSinkOps(); } statsTracker.flush(); @@ -376,7 +472,7 @@ public class PartitionWorker { long t1 = System.currentTimeMillis(); if (resultList == null || resultList.getData() == null || resultList.getData().isEmpty()) { - LOG.info("[PERF] {} read={}ms returned empty", entityType, t1 - t0); + LOG.debug("{} read={}ms returned empty", entityType, t1 - t0); return new BatchResult(0, 0, 0, null); } @@ -389,13 +485,22 @@ public class PartitionWorker { statsTracker.recordReaderBatch(readSuccessCount, readErrorCount, warningsCount); } + if (failureRecorder != null && readErrorCount > 0) { + for (EntityError entityError : listOrEmpty(resultList.getErrors())) { + String entityId = + entityError.getEntity() != null ? entityError.getEntity().toString() : null; + failureRecorder.recordReaderEntityFailure( + entityType, entityId, null, entityError.getMessage()); + } + } + Map contextData = createContextData(entityType, statsTracker); try { writeToSink(entityType, resultList, contextData); long t2 = System.currentTimeMillis(); - LOG.info( - "[PERF] {} read={}ms write={}ms total={}ms records={}", + LOG.debug( + "{} read={}ms write={}ms total={}ms records={}", entityType, t1 - t0, t2 - t1, @@ -429,8 +534,20 @@ public class PartitionWorker { PaginatedEntitiesSource source = new PaginatedEntitiesSource(entityType, limit, fields, 0); return source.readNextKeyset(keysetCursor); } else { + Long filterStartTs = null; + Long filterEndTs = null; + if (reindexConfig != null) { + long startTs = reindexConfig.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + filterStartTs = startTs; + filterEndTs = System.currentTimeMillis(); + } + } PaginatedEntityTimeSeriesSource source = - new PaginatedEntityTimeSeriesSource(entityType, limit, fields, 0); + (filterStartTs != null) + ? new PaginatedEntityTimeSeriesSource( + entityType, limit, fields, filterStartTs, filterEndTs) + : new PaginatedEntityTimeSeriesSource(entityType, limit, fields, 0); return source.readWithCursor(keysetCursor); } } @@ -442,7 +559,16 @@ public class PartitionWorker { if (!TIME_SERIES_ENTITIES.contains(entityType)) { int cursorOffset = (int) offset - 1; ListFilter filter = new ListFilter(Include.ALL); - return Entity.getEntityRepository(entityType).getCursorAtOffset(filter, cursorOffset); + String cursor = + Entity.getEntityRepository(entityType).getCursorAtOffset(filter, cursorOffset); + if (cursor == null) { + LOG.debug( + "getCursorAtOffset returned null for {} at offset {} (cursorOffset={})", + entityType, + offset, + cursorOffset); + } + return cursor; } else { return RestUtil.encodeCursor(String.valueOf(offset)); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java index bf316f2ffbc..cc9a3f8384d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java @@ -137,6 +137,38 @@ public class LoggingProgressListener implements ReindexingProgressListener { } } + @Override + public void onReaderFailure(String entityType, String entityId, String error, FailureType type) { + if (type == FailureType.ENTITY_NOT_FOUND) { + LOG.warn("Reader warning for {} [{}]: {}", entityType, entityId, error); + } else { + LOG.error("Reader failure for {} [{}] ({}): {}", entityType, entityId, type, error); + } + } + + @Override + public void onProcessFailure(String entityType, String entityId, String error) { + LOG.error("Process failure for {} [{}]: {}", entityType, entityId, error); + } + + @Override + public void onSinkFailure(String entityType, String entityId, String error) { + LOG.error("Sink failure for {} [{}]: {}", entityType, entityId, error); + } + + @Override + public void onSubIndexingCompleted(String entityType, String subIndex, StepStats subIndexStats) { + long success = + subIndexStats.getSuccessRecords() != null ? subIndexStats.getSuccessRecords() : 0; + long failed = subIndexStats.getFailedRecords() != null ? subIndexStats.getFailedRecords() : 0; + LOG.info( + "Sub-indexing completed for {} [{}] - Success: {}, Failed: {}", + entityType, + subIndex, + success, + failed); + } + @Override public void onJobStopped(Stats currentStats) { LOG.info("Reindexing job stopped by user request"); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListener.java index d34d12f452f..9c9fd0050fe 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListener.java @@ -6,6 +6,7 @@ import static org.openmetadata.service.socket.WebSocketManager.SEARCH_INDEX_JOB_ import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; import lombok.extern.slf4j.Slf4j; import org.openmetadata.schema.entity.app.App; import org.openmetadata.schema.entity.app.AppRunRecord; @@ -16,6 +17,7 @@ import org.openmetadata.schema.system.IndexingError; import org.openmetadata.schema.system.Stats; import org.openmetadata.schema.system.StepStats; import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.apps.bundles.searchIndex.QuartzOrchestratorContext; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingJobContext; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingProgressListener; @@ -30,19 +32,29 @@ import org.quartz.JobExecutionContext; public class QuartzProgressListener implements ReindexingProgressListener { private static final long WEBSOCKET_UPDATE_INTERVAL_MS = 2000; + private static final long DB_UPDATE_INTERVAL_MS = 5000; private static final int ERROR_THRESHOLD = 3; private final JobExecutionContext jobExecutionContext; private final EventPublisherJob jobData; private final App app; + private final Function jobRecordProvider; + private final QuartzOrchestratorContext.StatusPusher statusPusher; private volatile long lastWebSocketUpdate = 0; + private volatile long lastDbUpdate = 0; private final AtomicInteger pendingErrors = new AtomicInteger(0); public QuartzProgressListener( - JobExecutionContext jobExecutionContext, EventPublisherJob jobData, App app) { + JobExecutionContext jobExecutionContext, + EventPublisherJob jobData, + App app, + Function jobRecordProvider, + QuartzOrchestratorContext.StatusPusher statusPusher) { this.jobExecutionContext = jobExecutionContext; this.jobData = jobData; this.app = app; + this.jobRecordProvider = jobRecordProvider; + this.statusPusher = statusPusher; } @Override @@ -166,42 +178,81 @@ public class QuartzProgressListener implements ReindexingProgressListener { .getJobDataMap() .put(WEBSOCKET_STATUS_CHANNEL, SEARCH_INDEX_JOB_BROADCAST_CHANNEL); - updateRecordToDbAndNotify(); + updateRecordAndNotify(force); } catch (Exception ex) { - LOG.error("Failed to send updated stats with WebSocket", ex); + LOG.error("Failed to send updated stats", ex); } } - private void updateRecordToDbAndNotify() { - AppRunRecord appRecord = createAppRunRecord(); + private void updateRecordAndNotify(boolean forceDbUpdate) { + AppRunRecord appRecord = getUpdatedAppRunRecord(); + persistToDb(appRecord, forceDbUpdate); + broadcastViaWebSocket(appRecord); + } + + private void persistToDb(AppRunRecord appRecord, boolean force) { + if (statusPusher == null) { + return; + } + long currentTime = System.currentTimeMillis(); + if (!force && currentTime - lastDbUpdate < DB_UPDATE_INTERVAL_MS) { + return; + } + lastDbUpdate = currentTime; + try { + statusPusher.push(jobExecutionContext, appRecord, true); + } catch (Exception ex) { + LOG.error("Failed to persist app run record to database", ex); + } + } + + private void broadcastViaWebSocket(AppRunRecord appRecord) { if (WebSocketManager.getInstance() != null) { String messageJson = JsonUtils.pojoToJson(appRecord); WebSocketManager.getInstance() .broadCastMessageToAll(SEARCH_INDEX_JOB_BROADCAST_CHANNEL, messageJson); - LOG.debug("Broad-casted job updates via WebSocket. Status: {}", appRecord.getStatus()); } } - private AppRunRecord createAppRunRecord() { - AppRunRecord appRecord = new AppRunRecord(); - appRecord.setAppId(app != null ? app.getId() : null); - appRecord.setStartTime(jobData.getTimestamp()); + private AppRunRecord getUpdatedAppRunRecord() { + AppRunRecord appRecord = readExistingRecord(); appRecord.setStatus(AppRunRecord.Status.fromValue(jobData.getStatus().value())); + if (jobData.getStats() != null) { + SuccessContext ctx = appRecord.getSuccessContext(); + if (ctx == null) { + ctx = new SuccessContext(); + } + ctx.withAdditionalProperty("stats", jobData.getStats()); + appRecord.setSuccessContext(ctx); + } + if (jobData.getFailure() != null) { appRecord.setFailureContext( new FailureContext().withAdditionalProperty("failure", jobData.getFailure())); } - if (jobData.getStats() != null) { - appRecord.setSuccessContext( - new SuccessContext().withAdditionalProperty("stats", jobData.getStats())); - } - return appRecord; } + private AppRunRecord readExistingRecord() { + if (jobRecordProvider != null) { + try { + AppRunRecord existing = jobRecordProvider.apply(jobExecutionContext); + if (existing != null) { + return existing; + } + } catch (Exception ex) { + LOG.debug("Could not read existing job record from context", ex); + } + } + AppRunRecord fallback = new AppRunRecord(); + fallback.setAppId(app != null ? app.getId() : null); + fallback.setStartTime(jobData.getTimestamp()); + return fallback; + } + /** Get the current job data for external access */ public EventPublisherJob getJobData() { return jobData; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/stats/StageStatsTracker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/stats/StageStatsTracker.java index eea204fbbb3..f9809bf5d68 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/stats/StageStatsTracker.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/stats/StageStatsTracker.java @@ -4,6 +4,7 @@ import java.util.UUID; import java.util.concurrent.atomic.AtomicLong; import lombok.Getter; import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingMetrics; import org.openmetadata.service.jdbi3.CollectionDAO; /** @@ -113,8 +114,8 @@ public class StageStatsTracker { long deadline = System.currentTimeMillis() + timeoutMs; while (pendingSinkOps.get() > 0) { if (System.currentTimeMillis() >= deadline) { - LOG.warn( - "Timed out waiting for {} pending sink operations for job {} entity {}", + LOG.debug( + "Await cycle expired with {} pending sink operations for job {} entity {}", pendingSinkOps.get(), jobId, entityType); @@ -135,6 +136,24 @@ public class StageStatsTracker { return pendingSinkOps.get(); } + /** + * Reconcile any remaining pending sink operations by recording them as successful. This should + * only be called after the bulk processor has been flushed — at that point, submitted records are + * either written or would have been reported as failures through the error handler. Pending ops + * that remain are callbacks that didn't fire in time, not actual write failures. + */ + public void reconcilePendingSinkOps() { + long remaining = pendingSinkOps.getAndSet(0); + if (remaining > 0) { + sink.add((int) remaining, 0, 0); + LOG.info( + "Reconciled {} pending sink operations as successful for job {} entity {}", + remaining, + jobId, + entityType); + } + } + public void recordVector(StatsResult result) { vector.record(result); checkFlush(); @@ -184,6 +203,19 @@ public class StageStatsTracker { return; } + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + if (rSuccess > 0) metrics.recordStageSuccess("reader", entityType, rSuccess); + if (rFailed > 0) metrics.recordStageFailed("reader", entityType, rFailed); + if (rWarnings > 0) metrics.recordStageWarnings("reader", entityType, rWarnings); + if (pSuccess > 0) metrics.recordStageSuccess("process", entityType, pSuccess); + if (pFailed > 0) metrics.recordStageFailed("process", entityType, pFailed); + if (sSuccess > 0) metrics.recordStageSuccess("sink", entityType, sSuccess); + if (sFailed > 0) metrics.recordStageFailed("sink", entityType, sFailed); + if (vSuccess > 0) metrics.recordStageSuccess("vector", entityType, vSuccess); + if (vFailed > 0) metrics.recordStageFailed("vector", entityType, vFailed); + } + try { statsDAO.incrementStats( recordId.toString(), diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java index c6d4ceef591..22e4a4c1d4b 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java @@ -7507,6 +7507,27 @@ public interface CollectionDAO { connectionType = POSTGRES) void markStaleEntriesStopped(@Bind("appId") String appId); + @ConnectionAwareSqlUpdate( + value = + "UPDATE apps_extension_time_series SET json = JSON_SET(json, '$.status', 'stopped') WHERE appName=:appName AND JSON_UNQUOTE(JSON_EXTRACT(json, '$.status')) = 'running' AND extension = 'status'", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "UPDATE apps_extension_time_series SET json = jsonb_set(json, '{status}', '\"stopped\"') WHERE appName = :appName AND json->>'status' = 'running' AND extension = 'status'", + connectionType = POSTGRES) + void markStaleEntriesStoppedByName(@Bind("appName") String appName); + + @ConnectionAwareSqlUpdate( + value = + "UPDATE apps_extension_time_series SET json = JSON_SET(json, '$.status', 'stopped') WHERE appName=:appName AND JSON_UNQUOTE(JSON_EXTRACT(json, '$.status')) = 'running' AND extension = 'status' AND timestamp < :beforeTimestamp", + connectionType = MYSQL) + @ConnectionAwareSqlUpdate( + value = + "UPDATE apps_extension_time_series SET json = jsonb_set(json, '{status}', '\"stopped\"') WHERE appName = :appName AND json->>'status' = 'running' AND extension = 'status' AND timestamp < :beforeTimestamp", + connectionType = POSTGRES) + void markStaleEntriesStoppedBefore( + @Bind("appName") String appName, @Bind("beforeTimestamp") long beforeTimestamp); + @ConnectionAwareSqlUpdate( value = "UPDATE apps_extension_time_series SET json = JSON_SET(json, '$.status', 'failed') WHERE JSON_UNQUOTE(JSON_EXTRACT(json, '$.status')) = 'running' AND extension = 'status'", @@ -9644,10 +9665,13 @@ public interface CollectionDAO { @SqlQuery( "SELECT * FROM search_index_partition WHERE jobId = :jobId AND status = 'PROCESSING' " - + "AND assignedServer = :serverId ORDER BY claimedAt DESC LIMIT 1") + + "AND assignedServer = :serverId AND claimedAt = :claimedAt " + + "ORDER BY priority DESC, entityType, partitionIndex LIMIT 1") @RegisterRowMapper(SearchIndexPartitionMapper.class) SearchIndexPartitionRecord findLatestClaimedPartition( - @Bind("jobId") String jobId, @Bind("serverId") String serverId); + @Bind("jobId") String jobId, + @Bind("serverId") String serverId, + @Bind("claimedAt") long claimedAt); @SqlQuery( "SELECT * FROM search_index_partition WHERE jobId = :jobId AND status = :status " diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java index 53327617029..d120ed1a043 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java @@ -1224,6 +1224,11 @@ public abstract class EntityRepository { public String getCursorAtOffset(ListFilter filter, int offset) { List jsons = dao.listAfter(filter, 1, offset); if (jsons.isEmpty()) { + LOG.debug( + "getCursorAtOffset for {} at offset {} returned empty (filter condition={})", + entityType, + offset, + filter.getCondition(dao.getTableName())); return null; } T entity = JsonUtils.readValue(jsons.get(0), entityClass); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/HikariCPDataSourceFactory.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/HikariCPDataSourceFactory.java index 375d570cd66..c6723aa1bb9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/HikariCPDataSourceFactory.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/HikariCPDataSourceFactory.java @@ -45,7 +45,7 @@ public class HikariCPDataSourceFactory extends DataSourceFactory { private int minimumIdle = 10; @JsonProperty - @Max(100) + @Max(500) private int maximumPoolSize = 100; @JsonProperty private Long connectionTimeout; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/monitoring/MicrometerBundle.java b/openmetadata-service/src/main/java/org/openmetadata/service/monitoring/MicrometerBundle.java index a9658d13a34..2eaba5ef5a9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/monitoring/MicrometerBundle.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/monitoring/MicrometerBundle.java @@ -22,6 +22,7 @@ import java.io.IOException; import lombok.extern.slf4j.Slf4j; import org.glassfish.jersey.internal.inject.AbstractBinder; import org.openmetadata.service.OpenMetadataApplicationConfig; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingMetrics; /** * Dropwizard bundle for configuring Micrometer metrics with Prometheus backend. @@ -61,6 +62,8 @@ public class MicrometerBundle implements ConfiguredBundle diagnostics = new LinkedHashMap<>(); + diagnostics.put("timestamp", Instant.now().toString()); + diagnostics.put("jvm", collectJvmMetrics()); + diagnostics.put("jetty", collectJettyMetrics()); + diagnostics.put("database", collectDatabaseMetrics()); + diagnostics.put("bulk_executor", collectBulkExecutorMetrics()); + diagnostics.put("request_latency", collectRequestLatencyMetrics()); + return Response.ok(diagnostics).build(); + } + + private Map collectJvmMetrics() { + Map jvm = new LinkedHashMap<>(); + + MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean(); + MemoryUsage heapUsage = memoryBean.getHeapMemoryUsage(); + MemoryUsage nonHeapUsage = memoryBean.getNonHeapMemoryUsage(); + + long heapUsed = heapUsage.getUsed(); + long heapMax = heapUsage.getMax(); + jvm.put("heap_used_bytes", heapUsed); + jvm.put("heap_max_bytes", heapMax); + jvm.put("heap_usage_pct", heapMax > 0 ? Math.round(heapUsed * 1000.0 / heapMax) / 10.0 : 0.0); + jvm.put("non_heap_used_bytes", nonHeapUsage.getUsed()); + + long gcPauseTotalMs = 0; + long gcCount = 0; + for (GarbageCollectorMXBean gcBean : ManagementFactory.getGarbageCollectorMXBeans()) { + gcPauseTotalMs += gcBean.getCollectionTime(); + gcCount += gcBean.getCollectionCount(); + } + jvm.put("gc_pause_total_ms", gcPauseTotalMs); + jvm.put("gc_count", gcCount); + + ThreadMXBean threadBean = ManagementFactory.getThreadMXBean(); + jvm.put("thread_count", threadBean.getThreadCount()); + jvm.put("thread_peak", threadBean.getPeakThreadCount()); + + OperatingSystemMXBean osBean = ManagementFactory.getOperatingSystemMXBean(); + if (osBean instanceof com.sun.management.OperatingSystemMXBean sunOsBean) { + double processCpu = sunOsBean.getProcessCpuLoad(); + double systemCpu = sunOsBean.getCpuLoad(); + jvm.put("cpu_process_pct", Math.round(processCpu * 1000.0) / 10.0); + jvm.put("cpu_system_pct", Math.round(systemCpu * 1000.0) / 10.0); + } else { + jvm.put("cpu_process_pct", -1.0); + jvm.put("cpu_system_pct", -1.0); + } + + RuntimeMXBean runtimeBean = ManagementFactory.getRuntimeMXBean(); + jvm.put("uptime_seconds", runtimeBean.getUptime() / 1000); + + return jvm; + } + + private Map collectJettyMetrics() { + Map jetty = new LinkedHashMap<>(); + MeterRegistry registry = Metrics.globalRegistry; + + jetty.put("threads_current", gaugeValue(registry, "jetty.threads.current")); + jetty.put("threads_busy", gaugeValue(registry, "jetty.threads.busy")); + jetty.put("threads_idle", gaugeValue(registry, "jetty.threads.idle")); + jetty.put("threads_max", gaugeValue(registry, "jetty.threads.max")); + + double current = gaugeValue(registry, "jetty.threads.current"); + double busy = gaugeValue(registry, "jetty.threads.busy"); + jetty.put("utilization_pct", current > 0 ? Math.round(busy / current * 1000.0) / 10.0 : 0.0); + + jetty.put("queue_size", gaugeValue(registry, "jetty.queue.size")); + jetty.put("queue_time_avg_ms", gaugeValue(registry, "jetty.request.queue.time.ms")); + jetty.put("active_requests", gaugeValue(registry, "jetty.requests.active")); + + double virtualEnabled = gaugeValue(registry, "jetty.virtual.threads.enabled"); + jetty.put("virtual_threads_enabled", virtualEnabled > 0); + + return jetty; + } + + private Map collectDatabaseMetrics() { + Map db = new LinkedHashMap<>(); + MeterRegistry registry = Metrics.globalRegistry; + + double active = gaugeValue(registry, "hikaricp.connections.active"); + double idle = gaugeValue(registry, "hikaricp.connections.idle"); + double total = active + idle; + double max = gaugeValue(registry, "hikaricp.connections.max"); + double pending = gaugeValue(registry, "hikaricp.connections.pending"); + + db.put("pool_active", (int) active); + db.put("pool_idle", (int) idle); + db.put("pool_total", (int) total); + db.put("pool_max", (int) max); + db.put("pool_pending", (int) pending); + db.put("pool_usage_pct", max > 0 ? Math.round(active / max * 1000.0) / 10.0 : 0.0); + + double timeoutMs = gaugeValue(registry, "hikaricp.connections.timeout"); + db.put("connection_timeout_ms", timeoutMs > 0 ? (long) timeoutMs : 30000L); + + return db; + } + + private Map collectBulkExecutorMetrics() { + Map bulk = new LinkedHashMap<>(); + + try { + BulkExecutor executor = BulkExecutor.getInstance(); + int maxThreads = executor.getMaxThreads(); + int activeThreads = executor.getActiveCount(); + int queueDepth = executor.getQueueDepth(); + int queueCapacity = executor.getQueueSize(); + + bulk.put("max_threads", maxThreads); + bulk.put("active_threads", activeThreads); + bulk.put("queue_depth", queueDepth); + bulk.put("queue_capacity", queueCapacity); + bulk.put( + "queue_usage_pct", + queueCapacity > 0 ? Math.round(queueDepth * 1000.0 / queueCapacity) / 10.0 : 0.0); + bulk.put("has_capacity", executor.hasCapacity()); + } catch (Exception e) { + LOG.debug("Could not collect BulkExecutor metrics: {}", e.getMessage()); + bulk.put("error", "BulkExecutor not available"); + } + + return bulk; + } + + private Map collectRequestLatencyMetrics() { + Map latencyMap = new LinkedHashMap<>(); + MeterRegistry registry = Metrics.globalRegistry; + + Search totalTimerSearch = registry.find("request.latency.total"); + for (Meter meter : totalTimerSearch.meters()) { + if (!(meter instanceof Timer totalTimer)) { + continue; + } + + String endpoint = + totalTimer.getId().getTag("endpoint") != null + ? totalTimer.getId().getTag("endpoint") + : "unknown"; + String method = + totalTimer.getId().getTag("method") != null + ? totalTimer.getId().getTag("method") + : "unknown"; + String key = method + " " + endpoint; + + long count = totalTimer.count(); + if (count == 0) { + continue; + } + + double avgTotalMs = totalTimer.mean(TimeUnit.MILLISECONDS); + + double avgDbMs = timerMean(registry, "request.latency.database", endpoint, method); + double avgSearchMs = timerMean(registry, "request.latency.search", endpoint, method); + double avgInternalMs = timerMean(registry, "request.latency.internal", endpoint, method); + + double dbPct = avgTotalMs > 0 ? Math.round(avgDbMs / avgTotalMs * 1000.0) / 10.0 : 0.0; + double searchPct = + avgTotalMs > 0 ? Math.round(avgSearchMs / avgTotalMs * 1000.0) / 10.0 : 0.0; + double internalPct = + avgTotalMs > 0 ? Math.round(avgInternalMs / avgTotalMs * 1000.0) / 10.0 : 0.0; + + double avgDbOps = summaryMean(registry, "request.operations.database", endpoint, method); + double avgSearchOps = summaryMean(registry, "request.operations.search", endpoint, method); + + Map entry = new LinkedHashMap<>(); + entry.put("count", count); + entry.put("avg_total_ms", Math.round(avgTotalMs * 10.0) / 10.0); + entry.put("avg_db_ms", Math.round(avgDbMs * 10.0) / 10.0); + entry.put("avg_search_ms", Math.round(avgSearchMs * 10.0) / 10.0); + entry.put("avg_internal_ms", Math.round(avgInternalMs * 10.0) / 10.0); + entry.put("db_pct", dbPct); + entry.put("search_pct", searchPct); + entry.put("internal_pct", internalPct); + entry.put("avg_db_ops", Math.round(avgDbOps * 10.0) / 10.0); + entry.put("avg_search_ops", Math.round(avgSearchOps * 10.0) / 10.0); + + latencyMap.put(key, entry); + } + + return latencyMap; + } + + private static double gaugeValue(MeterRegistry registry, String name) { + Gauge gauge = registry.find(name).gauge(); + if (gauge != null) { + return gauge.value(); + } + return 0.0; + } + + private static double timerMean( + MeterRegistry registry, String name, String endpoint, String method) { + Timer timer = registry.find(name).tag("endpoint", endpoint).tag("method", method).timer(); + if (timer != null && timer.count() > 0) { + return timer.mean(TimeUnit.MILLISECONDS); + } + return 0.0; + } + + private static double summaryMean( + MeterRegistry registry, String name, String endpoint, String method) { + io.micrometer.core.instrument.DistributionSummary summary = + registry.find(name).tag("endpoint", endpoint).tag("method", method).summary(); + if (summary != null && summary.count() > 0) { + return summary.mean(); + } + return 0.0; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/DefaultRecreateHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/DefaultRecreateHandler.java index 914c21d8436..d3a3d5d74c6 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/DefaultRecreateHandler.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/DefaultRecreateHandler.java @@ -7,6 +7,7 @@ import java.util.Set; import lombok.extern.slf4j.Slf4j; import org.openmetadata.search.IndexMapping; import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.ReindexingMetrics; /** * Default implementation of RecreateHandler that provides zero-downtime index recreation. @@ -52,34 +53,60 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { if (canonicalIndex == null || stagedIndex == null) { LOG.error( - "Cannot finalize reindex for entity '{}'. Missing canonical or staged index name.", - entityType); + "Cannot finalize reindex for entity '{}'. canonicalIndex={}, stagedIndex={}", + entityType, + canonicalIndex, + stagedIndex); return; } - if (reindexSuccess) { + // Always-promote: partial data is better than no data. When reindex failed but the staged + // index has documents, promote it. Only delete if truly empty. + boolean shouldPromote = reindexSuccess; + if (!shouldPromote) { + long docCount = searchClient.getDocumentCount(stagedIndex); + if (docCount > 0) { + LOG.info( + "Reindex failed for entity '{}' but staged index '{}' has {} documents. " + + "Promoting partial data (partial data > no data).", + entityType, + stagedIndex, + docCount); + shouldPromote = true; + } else if (docCount == 0) { + LOG.info( + "Reindex failed for entity '{}' and staged index '{}' has 0 documents. " + + "Deleting empty staged index.", + entityType, + stagedIndex); + } else { + LOG.warn( + "Could not determine doc count for staged index '{}' (entity '{}'). " + + "Promoting to avoid data loss.", + stagedIndex, + entityType); + shouldPromote = true; + } + } + + if (shouldPromote) { try { Set aliasesToAttach = new HashSet<>(); - // Existing Aliases existingAliases.stream() .filter(alias -> alias != null && !alias.isBlank()) .forEach(aliasesToAttach::add); - // Canonical Alias if (!nullOrEmpty(canonicalAlias)) { aliasesToAttach.add(canonicalAlias); } - // Parent Aliases parentAliases.stream() .filter(alias -> alias != null && !alias.isBlank()) .forEach(aliasesToAttach::add); - // Remove any null or blank aliases aliasesToAttach.removeIf(alias -> alias == null || alias.isBlank()); - // Collect all old indices to delete (except staged) Set allEntityIndices = searchClient.listIndicesByPrefix(canonicalIndex); Set oldIndicesToDelete = new HashSet<>(); for (String oldIndex : allEntityIndices) { @@ -88,7 +115,13 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } } - // Canonical Indexes needs to be removed before attached that as aliases + LOG.debug( + "finalizeReindex entity '{}': aliases={}, oldIndices={}, stagedIndex={}", + entityType, + aliasesToAttach, + oldIndicesToDelete, + stagedIndex); + if (oldIndicesToDelete.contains(canonicalIndex)) { if (searchClient.indexExists(canonicalIndex)) { searchClient.deleteIndexWithBackoff(canonicalIndex); @@ -97,8 +130,6 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } } - // Atomically swap aliases from old indices to staged index - // This ensures zero-downtime: aliases point to new index before old ones are deleted if (!aliasesToAttach.isEmpty()) { boolean swapSuccess = searchClient.swapAliases(oldIndicesToDelete, stagedIndex, aliasesToAttach); @@ -111,12 +142,17 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } LOG.info( - "Promoted staged index '{}' to serve entity '{}' (aliases: {}).", + "Promoted staged index '{}' to serve entity '{}' (aliases: {}, reindexSuccess: {}).", stagedIndex, entityType, - aliasesToAttach); + aliasesToAttach, + reindexSuccess); + + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordPromotionSuccess(entityType); + } - // Delete old indices after successful alias swap (with backoff for snapshot scenarios) for (String oldIndex : oldIndicesToDelete) { try { if (searchClient.indexExists(oldIndex)) { @@ -131,6 +167,10 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } catch (Exception ex) { LOG.error( "Failed to promote staged index '{}' for entity '{}'.", stagedIndex, entityType, ex); + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.recordPromotionFailure(entityType); + } } } else { try { @@ -166,12 +206,39 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { if (canonicalIndex == null || stagedIndex == null) { LOG.error( - "Cannot promote index for entity '{}'. Missing canonical or staged index name.", - entityType); + "Cannot promote index for entity '{}'. canonicalIndex={}, stagedIndex={}", + entityType, + canonicalIndex, + stagedIndex); return; } - if (!reindexSuccess) { + // Always-promote: check doc count when reindex failed + boolean shouldPromote = reindexSuccess; + if (!shouldPromote) { + long docCount = searchClient.getDocumentCount(stagedIndex); + if (docCount > 0) { + LOG.info( + "Per-entity reindex failed for '{}' but staged index '{}' has {} documents. Promoting.", + entityType, + stagedIndex, + docCount); + shouldPromote = true; + } else if (docCount == 0) { + LOG.info( + "Per-entity reindex failed for '{}' and staged index '{}' is empty. Deleting.", + entityType, + stagedIndex); + } else { + LOG.warn( + "Could not determine doc count for staged index '{}' (entity '{}'). Promoting.", + stagedIndex, + entityType); + shouldPromote = true; + } + } + + if (!shouldPromote) { try { if (searchClient.indexExists(stagedIndex)) { searchClient.deleteIndexWithBackoff(stagedIndex); @@ -191,11 +258,9 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } try { - // Get aliases from indexMapping.json (not from old index) Set aliasesToAttach = getAliasesFromMapping(indexMapping, searchRepository.getClusterAlias()); - // Find old indices with this prefix (except staged) Set allEntityIndices = searchClient.listIndicesByPrefix(canonicalIndex); Set oldIndicesToDelete = new HashSet<>(); for (String oldIndex : allEntityIndices) { @@ -204,7 +269,13 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } } - // Canonical Indexes needs to be removed before attached that as aliases + LOG.debug( + "promoteEntityIndex '{}': aliases={}, oldIndices={}, stagedIndex={}", + entityType, + aliasesToAttach, + oldIndicesToDelete, + stagedIndex); + if (oldIndicesToDelete.contains(canonicalIndex)) { if (searchClient.indexExists(canonicalIndex)) { searchClient.deleteIndexWithBackoff(canonicalIndex); @@ -213,26 +284,35 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } } - // Atomically swap aliases from old indices to staged index - // This ensures zero-downtime: aliases point to new index before old ones are deleted if (!aliasesToAttach.isEmpty()) { boolean swapSuccess = searchClient.swapAliases(oldIndicesToDelete, stagedIndex, aliasesToAttach); if (!swapSuccess) { LOG.error( - "Failed to atomically swap aliases for entity '{}'. Old indices will not be deleted.", - entityType); + "Failed to atomically swap aliases for entity '{}'. " + + "oldIndices={}, stagedIndex={}, aliases={}", + entityType, + oldIndicesToDelete, + stagedIndex, + aliasesToAttach); return; } + } else { + LOG.warn("Entity '{}': aliasesToAttach is empty, skipping alias swap", entityType); } LOG.info( - "Promoted staged index '{}' to serve entity '{}' (aliases: {}).", + "Promoted staged index '{}' to serve entity '{}' (aliases: {}, reindexSuccess: {}).", stagedIndex, entityType, - aliasesToAttach); + aliasesToAttach, + reindexSuccess); + + ReindexingMetrics promoteMetrics = ReindexingMetrics.getInstance(); + if (promoteMetrics != null) { + promoteMetrics.recordPromotionSuccess(entityType); + } - // Delete old indices after successful alias swap (with backoff for snapshot scenarios) for (String oldIndex : oldIndicesToDelete) { try { if (searchClient.indexExists(oldIndex)) { @@ -247,6 +327,10 @@ public class DefaultRecreateHandler implements RecreateIndexHandler { } catch (Exception ex) { LOG.error( "Failed to promote staged index '{}' for entity '{}'.", stagedIndex, entityType, ex); + ReindexingMetrics promoteMetrics = ReindexingMetrics.getInstance(); + if (promoteMetrics != null) { + promoteMetrics.recordPromotionFailure(entityType); + } } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/IndexManagementClient.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/IndexManagementClient.java index 97ca657b7bd..4e8d8b99710 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/IndexManagementClient.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/IndexManagementClient.java @@ -143,4 +143,23 @@ public interface IndexManagementClient { Set aliases) {} List getAllIndexStats() throws IOException; + + /** + * Get the document count for a specific index. + * + * @param indexName the name of the index + * @return the number of documents in the index, or -1 if count cannot be determined + */ + default long getDocumentCount(String indexName) { + try { + for (IndexStats stats : getAllIndexStats()) { + if (stats.name().equals(indexName)) { + return stats.documents(); + } + } + } catch (Exception e) { + return -1; + } + return 0; + } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/PaginatedEntityTimeSeriesSource.java b/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/PaginatedEntityTimeSeriesSource.java index f341cc0b30f..afc3e0cd2f3 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/PaginatedEntityTimeSeriesSource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/PaginatedEntityTimeSeriesSource.java @@ -72,6 +72,21 @@ public class PaginatedEntityTimeSeriesSource this.endTs = endTs; } + public PaginatedEntityTimeSeriesSource( + String entityType, + int batchSize, + List fields, + int knownTotal, + Long startTs, + Long endTs) { + this.entityType = entityType; + this.batchSize = batchSize; + this.fields = fields; + this.stats.withTotalRecords(knownTotal).withSuccessRecords(0).withFailedRecords(0); + this.startTs = startTs; + this.endTs = endTs; + } + @Override public ResultList readNext(Map contextData) throws SearchIndexException { diff --git a/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json b/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json index 213c9b9c5b0..34a39128940 100644 --- a/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json +++ b/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json @@ -16,7 +16,8 @@ "initialBackoff": 1000, "maxBackoff": 10000, "searchIndexMappingLanguage": "EN", - "autoTune": false + "autoTune": false, + "timeSeriesMaxDays": 0 }, "appSchedule": { "scheduleTimeline": "Custom", diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java new file mode 100644 index 00000000000..5906b152bf2 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java @@ -0,0 +1,72 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("AdaptiveBackoff Tests") +class AdaptiveBackoffTest { + + @Test + @DisplayName("returns initial delay on first call") + void initialDelay() { + AdaptiveBackoff backoff = new AdaptiveBackoff(100, 2000); + assertEquals(100, backoff.nextDelay()); + } + + @Test + @DisplayName("doubles delay on each subsequent call") + void exponentialDoubling() { + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 10000); + assertEquals(50, backoff.nextDelay()); + assertEquals(100, backoff.nextDelay()); + assertEquals(200, backoff.nextDelay()); + assertEquals(400, backoff.nextDelay()); + assertEquals(800, backoff.nextDelay()); + } + + @Test + @DisplayName("caps at maxMs") + void capAtMax() { + AdaptiveBackoff backoff = new AdaptiveBackoff(100, 300); + assertEquals(100, backoff.nextDelay()); + assertEquals(200, backoff.nextDelay()); + assertEquals(300, backoff.nextDelay()); + assertEquals(300, backoff.nextDelay()); + } + + @Test + @DisplayName("reset returns to initial delay") + void resetToInitial() { + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 1000); + backoff.nextDelay(); + backoff.nextDelay(); + backoff.nextDelay(); + + backoff.reset(); + assertEquals(50, backoff.nextDelay()); + } + + @Test + @DisplayName("rejects invalid initialMs") + void rejectsInvalidInitialMs() { + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(0, 1000)); + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(-1, 1000)); + } + + @Test + @DisplayName("rejects maxMs less than initialMs") + void rejectsMaxLessThanInitial() { + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(200, 100)); + } + + @Test + @DisplayName("works when initialMs equals maxMs") + void initialEqualsMax() { + AdaptiveBackoff backoff = new AdaptiveBackoff(500, 500); + assertEquals(500, backoff.nextDelay()); + assertEquals(500, backoff.nextDelay()); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreakerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreakerTest.java new file mode 100644 index 00000000000..9833b6315aa --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/BulkCircuitBreakerTest.java @@ -0,0 +1,171 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("BulkCircuitBreaker Tests") +class BulkCircuitBreakerTest { + + private BulkCircuitBreaker breaker; + + @BeforeEach + void setUp() { + breaker = new BulkCircuitBreaker(3, 5000, 1000); + } + + @Test + @DisplayName("starts in CLOSED state") + void startsInClosedState() { + assertEquals(BulkCircuitBreaker.State.CLOSED, breaker.getState()); + assertTrue(breaker.allowRequest()); + } + + @Test + @DisplayName("transitions CLOSED → OPEN after threshold failures in window") + void closedToOpenOnThreshold() { + breaker.recordFailure(); + breaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.CLOSED, breaker.getState()); + + breaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.OPEN, breaker.getState()); + assertFalse(breaker.allowRequest()); + } + + @Test + @DisplayName("transitions OPEN → HALF_OPEN after probe interval") + void openToHalfOpenAfterInterval() { + BulkCircuitBreaker fastBreaker = new BulkCircuitBreaker(1, 5000, 50); + + fastBreaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.OPEN, fastBreaker.getState()); + assertFalse(fastBreaker.allowRequest()); + + try { + Thread.sleep(60); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + assertTrue(fastBreaker.allowRequest()); + assertEquals(BulkCircuitBreaker.State.HALF_OPEN, fastBreaker.getState()); + } + + @Test + @DisplayName("transitions HALF_OPEN → CLOSED on success") + void halfOpenToClosedOnSuccess() { + BulkCircuitBreaker fastBreaker = new BulkCircuitBreaker(1, 5000, 50); + + fastBreaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.OPEN, fastBreaker.getState()); + + try { + Thread.sleep(60); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + fastBreaker.allowRequest(); + assertEquals(BulkCircuitBreaker.State.HALF_OPEN, fastBreaker.getState()); + + fastBreaker.recordSuccess(); + assertEquals(BulkCircuitBreaker.State.CLOSED, fastBreaker.getState()); + assertTrue(fastBreaker.allowRequest()); + } + + @Test + @DisplayName("transitions HALF_OPEN → OPEN on failure") + void halfOpenToOpenOnFailure() { + BulkCircuitBreaker fastBreaker = new BulkCircuitBreaker(1, 5000, 50); + + fastBreaker.recordFailure(); + + try { + Thread.sleep(60); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + fastBreaker.allowRequest(); + assertEquals(BulkCircuitBreaker.State.HALF_OPEN, fastBreaker.getState()); + + fastBreaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.OPEN, fastBreaker.getState()); + } + + @Test + @DisplayName("failures outside window do not count toward threshold") + void expiryOutsideWindow() { + BulkCircuitBreaker shortWindow = new BulkCircuitBreaker(3, 100, 1000); + + shortWindow.recordFailure(); + shortWindow.recordFailure(); + + try { + Thread.sleep(150); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + + shortWindow.recordFailure(); + assertEquals(BulkCircuitBreaker.State.CLOSED, shortWindow.getState()); + } + + @Test + @DisplayName("reset forces CLOSED state") + void resetForcesClosed() { + breaker.recordFailure(); + breaker.recordFailure(); + breaker.recordFailure(); + assertEquals(BulkCircuitBreaker.State.OPEN, breaker.getState()); + + breaker.reset(); + assertEquals(BulkCircuitBreaker.State.CLOSED, breaker.getState()); + assertTrue(breaker.allowRequest()); + } + + @Test + @DisplayName("thread safety under concurrent access") + void threadSafety() throws InterruptedException { + BulkCircuitBreaker concurrentBreaker = new BulkCircuitBreaker(100, 10_000, 1000); + int threadCount = 10; + int failuresPerThread = 15; + CountDownLatch latch = new CountDownLatch(threadCount); + AtomicInteger allowedCount = new AtomicInteger(0); + + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + for (int i = 0; i < threadCount; i++) { + executor.submit( + () -> { + try { + for (int j = 0; j < failuresPerThread; j++) { + concurrentBreaker.recordFailure(); + if (concurrentBreaker.allowRequest()) { + allowedCount.incrementAndGet(); + } + } + } finally { + latch.countDown(); + } + }); + } + + latch.await(5, TimeUnit.SECONDS); + executor.shutdown(); + + BulkCircuitBreaker.State finalState = concurrentBreaker.getState(); + assertTrue( + finalState == BulkCircuitBreaker.State.OPEN + || finalState == BulkCircuitBreaker.State.CLOSED); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java new file mode 100644 index 00000000000..5f54e4f9f63 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java @@ -0,0 +1,67 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("EntityBatchSizeEstimator Tests") +class EntityBatchSizeEstimatorTest { + + @Test + @DisplayName("LARGE entities get smaller batch size") + void largeEntitiesGetSmallerBatch() { + int base = 200; + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("table", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("topic", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("dashboard", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("mlmodel", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("container", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("storedProcedure", base)); + } + + @Test + @DisplayName("LARGE entities respect minimum batch size of 25") + void largeEntitiesRespectMinimum() { + assertEquals(25, EntityBatchSizeEstimator.estimateBatchSize("table", 40)); + assertEquals(25, EntityBatchSizeEstimator.estimateBatchSize("table", 10)); + } + + @Test + @DisplayName("SMALL entities get larger batch size") + void smallEntitiesGetLargerBatch() { + int base = 200; + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("user", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("team", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("bot", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("role", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("policy", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("tag", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("classification", base)); + } + + @Test + @DisplayName("SMALL entities respect maximum batch size of 1000") + void smallEntitiesRespectMaximum() { + assertEquals(1000, EntityBatchSizeEstimator.estimateBatchSize("user", 600)); + assertEquals(1000, EntityBatchSizeEstimator.estimateBatchSize("user", 800)); + } + + @Test + @DisplayName("MEDIUM (unknown) entities get base batch size unchanged") + void mediumEntitiesUnchanged() { + int base = 200; + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("pipeline", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("database", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("glossaryTerm", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("unknownEntity", base)); + } + + @Test + @DisplayName("handles zero and negative base batch size gracefully") + void handlesZeroAndNegative() { + assertEquals(0, EntityBatchSizeEstimator.estimateBatchSize("table", 0)); + assertTrue(EntityBatchSizeEstimator.estimateBatchSize("table", -1) < 0); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriorityTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriorityTest.java new file mode 100644 index 00000000000..34d61f46962 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityPriorityTest.java @@ -0,0 +1,151 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("EntityPriority Tests") +class EntityPriorityTest { + + @Test + @DisplayName("Services sort before users/teams") + void servicesSortBeforeUsers() { + Set entities = Set.of("user", "databaseService", "team"); + List sorted = EntityPriority.sortByPriority(entities); + assertTrue(sorted.indexOf("databaseService") < sorted.indexOf("user")); + assertTrue(sorted.indexOf("databaseService") < sorted.indexOf("team")); + } + + @Test + @DisplayName("Users/teams sort before data assets") + void usersSortBeforeDataAssets() { + Set entities = Set.of("table", "user", "dashboard"); + List sorted = EntityPriority.sortByPriority(entities); + assertTrue(sorted.indexOf("user") < sorted.indexOf("table")); + assertTrue(sorted.indexOf("user") < sorted.indexOf("dashboard")); + } + + @Test + @DisplayName("Data assets sort before governance entities") + void dataAssetsSortBeforeGovernance() { + Set entities = Set.of("glossaryTerm", "table", "testCase"); + List sorted = EntityPriority.sortByPriority(entities); + assertTrue(sorted.indexOf("table") < sorted.indexOf("glossaryTerm")); + assertTrue(sorted.indexOf("table") < sorted.indexOf("testCase")); + } + + @Test + @DisplayName("Time series entities sort last") + void timeSeriesEntitiesSortLast() { + Set entities = Set.of("testCaseResult", "user", "table", "databaseService"); + List sorted = EntityPriority.sortByPriority(entities); + assertEquals("testCaseResult", sorted.get(sorted.size() - 1)); + } + + @Test + @DisplayName("Unknown entity types default to LOW tier") + void unknownEntitiesDefaultToLow() { + assertEquals(EntityPriority.Tier.LOW, EntityPriority.getTier("someUnknownEntity")); + } + + @Test + @DisplayName("Empty set returns empty list") + void emptySetReturnsEmptyList() { + List sorted = EntityPriority.sortByPriority(Set.of()); + assertTrue(sorted.isEmpty()); + } + + @Test + @DisplayName("Single entity returns single-element list") + void singleEntityReturnsSingleElement() { + List sorted = EntityPriority.sortByPriority(Set.of("table")); + assertEquals(1, sorted.size()); + assertEquals("table", sorted.get(0)); + } + + @Test + @DisplayName("Full priority ordering: CRITICAL > HIGH > MEDIUM > LOW > LOWEST") + void fullPriorityOrdering() { + Set entities = + Set.of("testCaseResult", "glossaryTerm", "table", "user", "databaseService"); + List sorted = EntityPriority.sortByPriority(entities); + + int serviceIdx = sorted.indexOf("databaseService"); + int userIdx = sorted.indexOf("user"); + int tableIdx = sorted.indexOf("table"); + int glossaryIdx = sorted.indexOf("glossaryTerm"); + int tsIdx = sorted.indexOf("testCaseResult"); + + assertTrue(serviceIdx < userIdx, "CRITICAL < HIGH"); + assertTrue(userIdx < tableIdx, "HIGH < MEDIUM"); + assertTrue(tableIdx < glossaryIdx, "MEDIUM < LOW"); + assertTrue(glossaryIdx < tsIdx, "LOW < LOWEST"); + } + + @Test + @DisplayName("Entities within same tier preserve relative order from input") + void sameTierPreservesOrder() { + LinkedHashSet entities = new LinkedHashSet<>(); + entities.add("dashboard"); + entities.add("table"); + entities.add("pipeline"); + List sorted = EntityPriority.sortByPriority(entities); + assertEquals(List.of("dashboard", "table", "pipeline"), sorted); + } + + @Test + @DisplayName("All service entities are CRITICAL tier") + void allServicesAreCritical() { + for (String svc : + List.of( + "databaseService", + "messagingService", + "dashboardService", + "pipelineService", + "mlmodelService", + "storageService", + "searchService", + "apiService", + "metadataService")) { + assertEquals(EntityPriority.Tier.CRITICAL, EntityPriority.getTier(svc), svc); + } + } + + @Test + @DisplayName("All time series entities are LOWEST tier") + void allTimeSeriesAreLowest() { + for (String ts : + List.of( + "entityReportData", + "rawCostAnalysisReportData", + "webAnalyticUserActivityReportData", + "webAnalyticEntityViewReportData", + "aggregatedCostAnalysisReportData", + "testCaseResolutionStatus", + "testCaseResult", + "queryCostRecord")) { + assertEquals(EntityPriority.Tier.LOWEST, EntityPriority.getTier(ts), ts); + } + } + + @Test + @DisplayName("Numeric priority maps correctly from tiers") + void numericPriorityMapsFromTiers() { + assertEquals(100, EntityPriority.getNumericPriority("databaseService")); + assertEquals(80, EntityPriority.getNumericPriority("user")); + assertEquals(60, EntityPriority.getNumericPriority("table")); + assertEquals(40, EntityPriority.getNumericPriority("glossaryTerm")); + assertEquals(20, EntityPriority.getNumericPriority("testCaseResult")); + } + + @Test + @DisplayName("Unknown entities get LOW numeric priority") + void unknownEntitiesGetLowNumericPriority() { + assertEquals(40, EntityPriority.getNumericPriority("someUnknownEntity")); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java new file mode 100644 index 00000000000..8a7fb871051 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java @@ -0,0 +1,108 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.service.exception.SearchIndexException; + +@DisplayName("EntityReader Retry Tests") +class EntityReaderRetryTest { + + @Test + @DisplayName("isTransientError detects timeout errors") + void detectsTimeoutErrors() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("Connection timeout while reading entities")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects connection errors") + void detectsConnectionErrors() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("java.net.ConnectException: Connection refused")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects pool exhaustion") + void detectsPoolExhaustion() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("Pool exhausted - no connections available")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects socket timeout") + void detectsSocketTimeout() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("java.net.SocketTimeoutException: Read timed out")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError returns false for non-transient errors") + void rejectsNonTransientErrors() { + SearchIndexException e = + new SearchIndexException(new IndexingError().withMessage("Entity not found: table.xyz")); + assertFalse(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError returns false for null message") + void handleNullMessage() { + SearchIndexException e = new SearchIndexException(new IndexingError()); + assertFalse(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("EntityReader constructor accepts custom retry configuration") + void customRetryConfiguration() { + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newSingleThreadExecutor(); + java.util.concurrent.atomic.AtomicBoolean stopped = + new java.util.concurrent.atomic.AtomicBoolean(false); + EntityReader reader = new EntityReader(executor, stopped, 5, 1000); + assertNotNull(reader); + executor.shutdown(); + } + + @Test + @DisplayName("EntityReader default constructor uses default retry values") + void defaultRetryConfiguration() { + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newSingleThreadExecutor(); + java.util.concurrent.atomic.AtomicBoolean stopped = + new java.util.concurrent.atomic.AtomicBoolean(false); + EntityReader reader = new EntityReader(executor, stopped); + assertNotNull(reader); + executor.shutdown(); + } + + @Test + @DisplayName("VectorCompletionResult.success creates completed result") + void vectorCompletionSuccess() { + VectorCompletionResult result = VectorCompletionResult.success(150); + assertTrue(result.completed()); + assertEquals(0, result.pendingTaskCount()); + assertEquals(150, result.waitedMillis()); + } + + @Test + @DisplayName("VectorCompletionResult.timeout creates timeout result") + void vectorCompletionTimeout() { + VectorCompletionResult result = VectorCompletionResult.timeout(5, 30000); + assertFalse(result.completed()); + assertEquals(5, result.pendingTaskCount()); + assertEquals(30000, result.waitedMillis()); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexErrorScenarioIntegrationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexErrorScenarioIntegrationTest.java new file mode 100644 index 00000000000..c3961149792 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexErrorScenarioIntegrationTest.java @@ -0,0 +1,785 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.system.EntityError; +import org.openmetadata.schema.type.Paging; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.search.IndexMapping; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.DistributedSearchIndexCoordinator; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.EntityCompletionTracker; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionStatus; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionWorker; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.PartitionWorker.PartitionResult; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexPartition; +import org.openmetadata.service.apps.bundles.searchIndex.distributed.ServerIdentityResolver; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.DefaultRecreateHandler; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.SearchClient; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.util.EntityUtil.Fields; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +@Slf4j +@DisplayName("Reindex Error Scenario Integration Tests") +class ReindexErrorScenarioIntegrationTest { + + @Mock private DistributedSearchIndexCoordinator coordinator; + @Mock private BulkSink bulkSink; + @Mock private CollectionDAO collectionDAO; + @Mock private CollectionDAO.SearchIndexServerStatsDAO statsDAO; + @Mock private CollectionDAO.SearchIndexFailureDAO failureDAO; + @Mock private EntityRepository mockRepository; + + private MockedStatic entityMock; + private MockedStatic serverIdMock; + private List capturedFailures; + private IndexingFailureRecorder failureRecorder; + private PartitionWorker worker; + private UUID jobId; + + @BeforeEach + void setUp() { + jobId = UUID.randomUUID(); + capturedFailures = new ArrayList<>(); + + entityMock = mockStatic(Entity.class); + serverIdMock = mockStatic(ServerIdentityResolver.class); + + entityMock.when(() -> Entity.getEntityRepository("table")).thenReturn(mockRepository); + entityMock + .when(() -> Entity.getFields(eq("table"), anyList())) + .thenReturn(new Fields(Collections.emptySet())); + + ServerIdentityResolver mockResolver = mock(ServerIdentityResolver.class); + when(mockResolver.getServerId()).thenReturn("test-server"); + serverIdMock.when(ServerIdentityResolver::getInstance).thenReturn(mockResolver); + + when(coordinator.getCollectionDAO()).thenReturn(collectionDAO); + when(collectionDAO.searchIndexServerStatsDAO()).thenReturn(statsDAO); + when(collectionDAO.searchIndexFailureDAO()).thenReturn(failureDAO); + + doAnswer( + invocation -> { + @SuppressWarnings("unchecked") + List records = + invocation.getArgument(0); + capturedFailures.addAll(records); + return null; + }) + .when(failureDAO) + .insertBatch(anyList()); + + lenient() + .when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())) + .thenReturn("encoded-cursor"); + + lenient().when(bulkSink.flushAndAwait(anyInt())).thenReturn(true); + lenient().when(bulkSink.getPendingVectorTaskCount()).thenReturn(0); + + failureRecorder = + new IndexingFailureRecorder(collectionDAO, jobId.toString(), "test-server", 100); + + worker = new PartitionWorker(coordinator, bulkSink, 100, null, false, failureRecorder, null); + } + + @AfterEach + void tearDown() { + if (failureRecorder != null) { + failureRecorder.close(); + } + if (entityMock != null) { + entityMock.close(); + } + if (serverIdMock != null) { + serverIdMock.close(); + } + } + + private SearchIndexPartition createPartition(long rangeStart, long rangeEnd) { + return SearchIndexPartition.builder() + .id(UUID.randomUUID()) + .jobId(jobId) + .entityType("table") + .partitionIndex(0) + .rangeStart(rangeStart) + .rangeEnd(rangeEnd) + .estimatedCount(rangeEnd - rangeStart) + .workUnits(rangeEnd - rangeStart) + .priority(0) + .status(PartitionStatus.PENDING) + .cursor(0L) + .processedCount(0L) + .successCount(0L) + .failedCount(0L) + .claimableAt(System.currentTimeMillis()) + .build(); + } + + private ResultList createResultList(int count, String nextCursor) { + List entities = new ArrayList<>(); + for (int i = 0; i < count; i++) { + entities.add(mock(EntityInterface.class)); + } + Paging paging = new Paging(); + paging.setAfter(nextCursor); + paging.setTotal(count); + ResultList result = new ResultList<>(entities); + result.setPaging(paging); + return result; + } + + private ResultList createResultListWithErrors( + int successCount, int errorCount, String nextCursor) { + List entities = new ArrayList<>(); + for (int i = 0; i < successCount; i++) { + entities.add(mock(EntityInterface.class)); + } + List errors = new ArrayList<>(); + for (int i = 0; i < errorCount; i++) { + errors.add(new EntityError().withMessage("Error reading entity " + i).withEntity("eid-" + i)); + } + Paging paging = new Paging(); + paging.setAfter(nextCursor); + paging.setTotal(successCount + errorCount); + ResultList result = new ResultList<>(entities); + result.setPaging(paging); + result.setErrors(errors); + return result; + } + + private void stubListAfterKeysetThrowFirst(Throwable t) { + doThrow(t) + .when(mockRepository) + .listAfterKeyset( + any(ListFilter.class), anyInt(), any(), anyInt(), eq(true), any(Fields.class)); + } + + private void stubListAfterKeysetViaAnswer(List responses) { + AtomicInteger callIndex = new AtomicInteger(0); + doAnswer( + invocation -> { + int idx = callIndex.getAndIncrement(); + if (idx < responses.size()) { + Object resp = responses.get(idx); + if (resp instanceof Throwable t) { + throw t; + } + return resp; + } + return createResultList(0, null); + }) + .when(mockRepository) + .listAfterKeyset( + any(ListFilter.class), anyInt(), any(), anyInt(), eq(true), any(Fields.class)); + } + + @Nested + @DisplayName("1. Reader Failure Tests") + class ReaderFailureTests { + + @Test + @DisplayName("Reader throws mid-partition — batches 1,3 OK, batch 2 throws") + void testReaderThrowsMidPartition() { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch3 = createResultList(100, null); + + stubListAfterKeysetViaAnswer( + List.of(batch1, new RuntimeException("DB connection lost"), batch3)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(200, result.successCount()); + assertEquals(100, result.failedCount()); + assertEquals(100, result.readerFailed()); + assertFalse(result.wasStopped()); + verify(coordinator).completePartition(eq(partition.getId()), eq(200L), eq(100L)); + } + + @Test + @DisplayName("Reader returns empty — data exhausted early") + void testReaderReturnsEmpty() { + SearchIndexPartition partition = createPartition(0, 200); + + ResultList batch1 = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch1)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(100, result.successCount()); + assertEquals(0, result.failedCount()); + assertFalse(result.wasStopped()); + } + + @Test + @DisplayName("Reader throws on first batch") + void testReaderThrowsOnFirstBatch() { + SearchIndexPartition partition = createPartition(0, 100); + + stubListAfterKeysetThrowFirst(new RuntimeException("Table not found")); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(0, result.successCount()); + assertEquals(100, result.failedCount()); + assertEquals(100, result.readerFailed()); + } + + @Test + @DisplayName("Reader returns ResultList with EntityErrors") + void testReaderReturnsEntityErrors() { + SearchIndexPartition partition = createPartition(0, 100); + + ResultList batchWithErrors = createResultListWithErrors(95, 5, null); + stubListAfterKeysetViaAnswer(List.of(batchWithErrors)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(95, result.successCount()); + assertEquals(5, result.failedCount()); + + failureRecorder.flush(); + long readerFailureCount = + capturedFailures.stream().filter(r -> "READER".equals(r.getFailureStage())).count(); + assertEquals(5, readerFailureCount); + } + + @Test + @DisplayName("Reader throws on last batch only") + void testReaderThrowsOnLastBatch() { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch2 = createResultList(100, "cursor-2"); + + stubListAfterKeysetViaAnswer(List.of(batch1, batch2, new RuntimeException("Timeout"))); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(200, result.successCount()); + assertEquals(100, result.failedCount()); + } + } + + @Nested + @DisplayName("2. Sink Failure Tests") + class SinkFailureTests { + + @Test + @DisplayName("Sink throws on write") + void testSinkThrowsOnWrite() throws Exception { + SearchIndexPartition partition = createPartition(0, 100); + + ResultList batch1 = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch1)); + + doThrow(new RuntimeException("Connection reset")) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(0, result.successCount()); + assertEquals(100, result.failedCount()); + verify(coordinator).completePartition(eq(partition.getId()), eq(0L), eq(100L)); + } + + @Test + @DisplayName("Sink fails second batch only") + void testSinkFailsSecondBatchOnly() throws Exception { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch2 = createResultList(100, "cursor-2"); + ResultList batch3 = createResultList(100, null); + + stubListAfterKeysetViaAnswer(List.of(batch1, batch2, batch3)); + + AtomicInteger writeCallCount = new AtomicInteger(0); + doAnswer( + invocation -> { + int call = writeCallCount.incrementAndGet(); + if (call == 2) { + throw new RuntimeException("Connection reset"); + } + return null; + }) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(200, result.successCount()); + assertEquals(100, result.failedCount()); + } + + @Test + @DisplayName("Sink fails all batches") + void testSinkFailsAllBatches() throws Exception { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch2 = createResultList(100, "cursor-2"); + ResultList batch3 = createResultList(100, null); + + stubListAfterKeysetViaAnswer(List.of(batch1, batch2, batch3)); + + doThrow(new RuntimeException("Connection reset")) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(0, result.successCount()); + assertEquals(300, result.failedCount()); + } + } + + @Nested + @DisplayName("3. Process Failure Tests") + class ProcessFailureTests { + + @Test + @DisplayName("Sink write throws RuntimeException — treated as SINK error") + void testDocBuildFailureTreatedAsSink() throws Exception { + SearchIndexPartition partition = createPartition(0, 100); + + ResultList batch1 = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch1)); + + doThrow(new RuntimeException("Failed to serialize")) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(0, result.successCount()); + assertEquals(100, result.failedCount()); + assertEquals(0, result.readerFailed()); + + failureRecorder.flush(); + long sinkFailures = + capturedFailures.stream().filter(r -> "SINK".equals(r.getFailureStage())).count(); + assertTrue(sinkFailures > 0); + } + + @Test + @DisplayName("Fatal exception in updatePartitionProgress — failPartition called") + void testFatalExceptionCallsFailPartition() { + SearchIndexPartition partition = createPartition(0, 100); + + doThrow(new NullPointerException("Unexpected null")) + .when(coordinator) + .updatePartitionProgress(any(SearchIndexPartition.class)); + + worker.processPartition(partition); + + verify(coordinator).failPartition(eq(partition.getId()), anyString()); + } + } + + @Nested + @DisplayName("4. Vector Embedding Failure Tests") + class VectorEmbeddingFailureTests { + + @Test + @DisplayName("Vector timeout — partition completes normally with warning") + void testVectorTimeout() { + SearchIndexPartition partition = createPartition(0, 100); + + ResultList batch1 = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch1)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + when(bulkSink.getPendingVectorTaskCount()).thenReturn(5); + when(bulkSink.awaitVectorCompletion(120)).thenReturn(false); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(100, result.successCount()); + assertFalse(result.wasStopped()); + verify(coordinator).completePartition(eq(partition.getId()), eq(100L), eq(0L)); + } + + @Test + @DisplayName("Vector tasks complete normally") + void testVectorTasksCompleteNormally() { + SearchIndexPartition partition = createPartition(0, 100); + + ResultList batch1 = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch1)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + when(bulkSink.getPendingVectorTaskCount()).thenReturn(3).thenReturn(0); + when(bulkSink.awaitVectorCompletion(120)).thenReturn(true); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(100, result.successCount()); + verify(coordinator).completePartition(eq(partition.getId()), eq(100L), eq(0L)); + } + } + + @Nested + @DisplayName("5. Promotion Failure Tests") + class PromotionFailureTests { + + private SearchClient setupPromotionMocks(SearchRepository searchRepo) { + SearchClient searchClient = mock(SearchClient.class); + when(searchRepo.getSearchClient()).thenReturn(searchClient); + when(searchRepo.getClusterAlias()).thenReturn(""); + + IndexMapping indexMapping = + IndexMapping.builder() + .indexName("table_search_index") + .alias("table") + .parentAliases(List.of("all")) + .childAliases(List.of()) + .build(); + when(searchRepo.getIndexMapping("table")).thenReturn(indexMapping); + + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepo); + + return searchClient; + } + + @Test + @DisplayName("swapAliases returns false — old indices NOT deleted") + void testSwapAliasesReturnsFalse() { + SearchRepository searchRepo = mock(SearchRepository.class); + SearchClient searchClient = setupPromotionMocks(searchRepo); + + when(searchClient.listIndicesByPrefix("table_search_index")) + .thenReturn(Set.of("table_search_index_rebuild_old", "table_search_index_rebuild_new")); + when(searchClient.indexExists(anyString())).thenReturn(false); + when(searchClient.swapAliases(any(), anyString(), any())).thenReturn(false); + + EntityReindexContext context = + EntityReindexContext.builder() + .entityType("table") + .canonicalIndex("table_search_index") + .stagedIndex("table_search_index_rebuild_new") + .build(); + + new DefaultRecreateHandler().promoteEntityIndex(context, true); + + verify(searchClient, never()).deleteIndexWithBackoff(eq("table_search_index_rebuild_old")); + } + + @Test + @DisplayName("getDocumentCount returns -1 — should promote to avoid data loss") + void testDocCountUnknownPromotes() { + SearchRepository searchRepo = mock(SearchRepository.class); + SearchClient searchClient = setupPromotionMocks(searchRepo); + + when(searchClient.getDocumentCount("table_search_index_rebuild_new")).thenReturn(-1L); + when(searchClient.listIndicesByPrefix("table_search_index")).thenReturn(Set.of()); + when(searchClient.indexExists(anyString())).thenReturn(false); + when(searchClient.swapAliases(any(), anyString(), any())).thenReturn(true); + + EntityReindexContext context = + EntityReindexContext.builder() + .entityType("table") + .canonicalIndex("table_search_index") + .stagedIndex("table_search_index_rebuild_new") + .build(); + + new DefaultRecreateHandler().promoteEntityIndex(context, false); + + verify(searchClient).swapAliases(any(), eq("table_search_index_rebuild_new"), any()); + } + + @Test + @DisplayName("Promotion callback throws — entity still in promotedEntities") + void testPromotionCallbackThrowsEntityStillPromoted() { + EntityCompletionTracker tracker = new EntityCompletionTracker(jobId); + tracker.initializeEntity("table", 1); + tracker.setOnEntityComplete( + (entityType, success) -> { + throw new RuntimeException("Promotion failed"); + }); + + tracker.recordPartitionComplete("table", false); + + assertTrue(tracker.isPromoted("table")); + } + + @Test + @DisplayName("Zero-doc entity, reindex failed — should NOT promote") + void testZeroDocReindexFailedNoPromotion() { + SearchRepository searchRepo = mock(SearchRepository.class); + SearchClient searchClient = setupPromotionMocks(searchRepo); + + when(searchClient.getDocumentCount("table_search_index_rebuild_new")).thenReturn(0L); + when(searchClient.indexExists("table_search_index_rebuild_new")).thenReturn(true); + + EntityReindexContext context = + EntityReindexContext.builder() + .entityType("table") + .canonicalIndex("table_search_index") + .stagedIndex("table_search_index_rebuild_new") + .build(); + + new DefaultRecreateHandler().promoteEntityIndex(context, false); + + verify(searchClient, never()).swapAliases(any(), anyString(), any()); + verify(searchClient).deleteIndexWithBackoff("table_search_index_rebuild_new"); + } + + @Test + @DisplayName("Zero-doc entity, reindex succeeded — should promote") + void testZeroDocReindexSuccessPromotes() { + SearchRepository searchRepo = mock(SearchRepository.class); + SearchClient searchClient = setupPromotionMocks(searchRepo); + + when(searchClient.listIndicesByPrefix("table_search_index")).thenReturn(Set.of()); + when(searchClient.indexExists(anyString())).thenReturn(false); + when(searchClient.swapAliases(any(), anyString(), any())).thenReturn(true); + + EntityReindexContext context = + EntityReindexContext.builder() + .entityType("table") + .canonicalIndex("table_search_index") + .stagedIndex("table_search_index_rebuild_new") + .build(); + + new DefaultRecreateHandler().promoteEntityIndex(context, true); + + verify(searchClient).swapAliases(any(), eq("table_search_index_rebuild_new"), any()); + } + } + + @Nested + @DisplayName("6. Mixed Failure Tests") + class MixedFailureTests { + + @Test + @DisplayName("Reader + sink failures in same partition") + void testReaderAndSinkFailuresSamePartition() throws Exception { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch2 = createResultList(100, "cursor-2"); + ResultList batch3 = createResultList(100, null); + + stubListAfterKeysetViaAnswer(List.of(new RuntimeException("DB error"), batch2, batch3)); + + AtomicInteger writeCallCount = new AtomicInteger(0); + doAnswer( + invocation -> { + int call = writeCallCount.incrementAndGet(); + if (call == 2) { + throw new RuntimeException("Sink error"); + } + return null; + }) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(100, result.successCount()); + assertEquals(200, result.failedCount()); + assertEquals(100, result.readerFailed()); + } + + @Test + @DisplayName("Multiple processPartition calls have independent stats") + void testMultipleCallsIndependentStats() { + ResultList batch = createResultList(100, null); + stubListAfterKeysetViaAnswer(List.of(batch)); + + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + SearchIndexPartition partition1 = createPartition(0, 100); + PartitionResult result1 = worker.processPartition(partition1); + assertEquals(100, result1.successCount()); + assertEquals(0, result1.failedCount()); + + stubListAfterKeysetThrowFirst(new RuntimeException("Failure")); + when(mockRepository.getCursorAtOffset(any(ListFilter.class), anyInt())).thenReturn(null); + + SearchIndexPartition partition2 = createPartition(0, 100); + PartitionResult result2 = worker.processPartition(partition2); + assertEquals(0, result2.successCount()); + assertEquals(100, result2.failedCount()); + } + } + + @Nested + @DisplayName("7. Stats Accuracy Tests") + class StatsAccuracyTests { + + @Test + @DisplayName("Stats consistent after reader failure: success + failed == total") + void testStatsConsistentAfterReaderFailure() { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch3 = createResultList(100, null); + + stubListAfterKeysetViaAnswer(List.of(batch1, new RuntimeException("DB error"), batch3)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(300, result.successCount() + result.failedCount()); + } + + @Test + @DisplayName("Stats consistent after sink failure: success + failed == total") + void testStatsConsistentAfterSinkFailure() throws Exception { + SearchIndexPartition partition = createPartition(0, 200); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch2 = createResultList(100, null); + + stubListAfterKeysetViaAnswer(List.of(batch1, batch2)); + + AtomicInteger writeCallCount = new AtomicInteger(0); + doAnswer( + invocation -> { + int call = writeCallCount.incrementAndGet(); + if (call == 1) { + throw new RuntimeException("Sink error"); + } + return null; + }) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(200, result.successCount() + result.failedCount()); + } + + @Test + @DisplayName("Stats consistent after mixed failures: success + failed == total") + void testStatsConsistentAfterMixedFailures() throws Exception { + SearchIndexPartition partition = createPartition(0, 500); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch3 = createResultList(100, "cursor-3"); + ResultList batch4 = createResultList(100, "cursor-4"); + ResultList batch5 = createResultList(100, null); + + stubListAfterKeysetViaAnswer( + List.of(batch1, new RuntimeException("Reader error"), batch3, batch4, batch5)); + + AtomicInteger writeCallCount = new AtomicInteger(0); + doAnswer( + invocation -> { + int call = writeCallCount.incrementAndGet(); + if (call == 3) { + throw new RuntimeException("Sink error"); + } + return null; + }) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertEquals(500, result.successCount() + result.failedCount()); + } + } + + @Nested + @DisplayName("8. Partition Lifecycle Tests") + class PartitionLifecycleTests { + + @Test + @DisplayName("Worker stopped mid-partition — wasStopped=true, completePartition NOT called") + void testWorkerStoppedMidPartition() throws Exception { + SearchIndexPartition partition = createPartition(0, 300); + + ResultList batch1 = createResultList(100, "cursor-1"); + ResultList batch2 = createResultList(100, "cursor-2"); + + stubListAfterKeysetViaAnswer(List.of(batch1, batch2)); + + doAnswer( + invocation -> { + worker.stop(); + return null; + }) + .when(bulkSink) + .write(anyList(), any(Map.class)); + + PartitionResult result = worker.processPartition(partition); + + assertTrue(result.wasStopped()); + verify(coordinator, never()).completePartition(any(UUID.class), anyLong(), anyLong()); + } + + @Test + @DisplayName("Coordinator fails partition on fatal error") + void testCoordinatorFailsPartitionOnFatalError() { + SearchIndexPartition partition = createPartition(0, 100); + + doThrow(new NullPointerException("Unexpected")) + .when(coordinator) + .updatePartitionProgress(any(SearchIndexPartition.class)); + + worker.processPartition(partition); + + verify(coordinator).failPartition(eq(partition.getId()), anyString()); + verify(coordinator, never()).completePartition(any(UUID.class), anyLong(), anyLong()); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTimeSeriesTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTimeSeriesTest.java new file mode 100644 index 00000000000..e94c0be3b16 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTimeSeriesTest.java @@ -0,0 +1,137 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.system.EventPublisherJob; + +@DisplayName("ReindexingConfiguration Time Series Tests") +class ReindexingConfigurationTimeSeriesTest { + + @Test + @DisplayName("getTimeSeriesStartTs returns correct timestamp for default days") + void defaultDaysReturnsCorrectTimestamp() { + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("testCaseResult")) + .timeSeriesMaxDays(15) + .build(); + + long startTs = config.getTimeSeriesStartTs("testCaseResult"); + long expectedApprox = System.currentTimeMillis() - (15 * 86_400_000L); + assertTrue( + Math.abs(startTs - expectedApprox) < 1000, + "Start timestamp should be approximately 15 days ago"); + } + + @Test + @DisplayName("getTimeSeriesStartTs uses entity-specific override when configured") + void entitySpecificOverride() { + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("testCaseResult", "entityReportData")) + .timeSeriesMaxDays(15) + .timeSeriesEntityDays(Map.of("testCaseResult", 30)) + .build(); + + long testCaseStartTs = config.getTimeSeriesStartTs("testCaseResult"); + long reportDataStartTs = config.getTimeSeriesStartTs("entityReportData"); + + long expected30Days = System.currentTimeMillis() - (30 * 86_400_000L); + long expected15Days = System.currentTimeMillis() - (15 * 86_400_000L); + + assertTrue( + Math.abs(testCaseStartTs - expected30Days) < 1000, + "testCaseResult should use 30-day override"); + assertTrue( + Math.abs(reportDataStartTs - expected15Days) < 1000, + "entityReportData should use default 15 days"); + } + + @Test + @DisplayName("getTimeSeriesStartTs returns -1 when days is 0 (no filtering)") + void zeroDaysReturnsNegativeOne() { + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("testCaseResult")) + .timeSeriesMaxDays(0) + .build(); + + assertEquals(-1, config.getTimeSeriesStartTs("testCaseResult")); + } + + @Test + @DisplayName("getTimeSeriesStartTs returns -1 when days is -1 (no filtering)") + void negativeDaysReturnsNegativeOne() { + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("testCaseResult")) + .timeSeriesMaxDays(-1) + .build(); + + assertEquals(-1, config.getTimeSeriesStartTs("testCaseResult")); + } + + @Test + @DisplayName("getTimeSeriesStartTs returns -1 when entity override is 0") + void entityOverrideZeroReturnsNegativeOne() { + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("testCaseResult")) + .timeSeriesMaxDays(15) + .timeSeriesEntityDays(Map.of("testCaseResult", 0)) + .build(); + + assertEquals(-1, config.getTimeSeriesStartTs("testCaseResult")); + } + + @Test + @DisplayName("from(EventPublisherJob) correctly reads new fields") + void fromJobReadsNewFields() { + EventPublisherJob job = new EventPublisherJob(); + job.setTimeSeriesMaxDays(30); + job.setTimeSeriesEntityDays(Map.of("testCaseResult", 60)); + job.setEntities(Set.of("table")); + + ReindexingConfiguration config = ReindexingConfiguration.from(job); + + assertEquals(30, config.timeSeriesMaxDays()); + assertEquals(Map.of("testCaseResult", 60), config.timeSeriesEntityDays()); + } + + @Test + @DisplayName("from(EventPublisherJob) uses defaults when fields are null") + void fromJobUsesDefaults() { + EventPublisherJob job = new EventPublisherJob(); + job.setEntities(Set.of("table")); + + ReindexingConfiguration config = ReindexingConfiguration.from(job); + + assertEquals(0, config.timeSeriesMaxDays()); + assertEquals(Collections.emptyMap(), config.timeSeriesEntityDays()); + } + + @Test + @DisplayName("Builder propagates new fields") + void builderPropagatesFields() { + Map entityDays = Map.of("queryCostRecord", 7, "testCaseResult", 45); + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of("table")) + .timeSeriesMaxDays(20) + .timeSeriesEntityDays(entityDays) + .build(); + + assertEquals(20, config.timeSeriesMaxDays()); + assertEquals(entityDays, config.timeSeriesEntityDays()); + + long queryCostStartTs = config.getTimeSeriesStartTs("queryCostRecord"); + long expected7Days = System.currentTimeMillis() - (7 * 86_400_000L); + assertTrue(Math.abs(queryCostStartTs - expected7Days) < 1000); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetricsTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetricsTest.java new file mode 100644 index 00000000000..ac3e19eb5d6 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingMetricsTest.java @@ -0,0 +1,445 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.lang.reflect.Field; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +@DisplayName("ReindexingMetrics Tests") +class ReindexingMetricsTest { + + private SimpleMeterRegistry meterRegistry; + + @BeforeEach + void setUp() throws Exception { + resetSingleton(); + meterRegistry = new SimpleMeterRegistry(); + ReindexingMetrics.initialize(meterRegistry); + } + + private void resetSingleton() throws Exception { + Field instanceField = ReindexingMetrics.class.getDeclaredField("instance"); + instanceField.setAccessible(true); + instanceField.set(null, null); + } + + @Nested + @DisplayName("Initialization") + class InitializationTests { + + @Test + @DisplayName("getInstance returns non-null after initialize") + void testGetInstanceAfterInitialize() { + assertNotNull(ReindexingMetrics.getInstance()); + } + + @Test + @DisplayName("getInstance returns null before initialize") + void testGetInstanceBeforeInitialize() throws Exception { + resetSingleton(); + assertNull(ReindexingMetrics.getInstance()); + } + + @Test + @DisplayName("Double initialization is idempotent") + void testDoubleInitialize() { + ReindexingMetrics first = ReindexingMetrics.getInstance(); + SimpleMeterRegistry secondRegistry = new SimpleMeterRegistry(); + ReindexingMetrics.initialize(secondRegistry); + assertEquals(first, ReindexingMetrics.getInstance()); + } + + @Test + @DisplayName("All metrics are registered") + void testAllMetricsRegistered() { + assertNotNull(meterRegistry.find("reindexing.jobs").tag("status", "started").counter()); + assertNotNull(meterRegistry.find("reindexing.jobs").tag("status", "completed").counter()); + assertNotNull(meterRegistry.find("reindexing.jobs").tag("status", "failed").counter()); + assertNotNull(meterRegistry.find("reindexing.jobs").tag("status", "stopped").counter()); + + assertNotNull( + meterRegistry.find("reindexing.job.duration").tag("status", "completed").timer()); + assertNotNull(meterRegistry.find("reindexing.job.duration").tag("status", "failed").timer()); + assertNotNull(meterRegistry.find("reindexing.job.duration").tag("status", "stopped").timer()); + + assertNotNull(meterRegistry.find("reindexing.jobs.active").gauge()); + + assertNotNull(meterRegistry.find("reindexing.bulk.duration").tag("success", "true").timer()); + assertNotNull(meterRegistry.find("reindexing.bulk.duration").tag("success", "false").timer()); + assertNotNull(meterRegistry.find("reindexing.bulk.payload.size").summary()); + assertNotNull(meterRegistry.find("reindexing.sink.pending").gauge()); + assertNotNull(meterRegistry.find("reindexing.backpressure.events").counter()); + } + } + + @Nested + @DisplayName("Job Lifecycle") + class JobLifecycleTests { + + @Test + @DisplayName("recordJobStarted increments counter and active gauge") + void testRecordJobStarted() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + + assertEquals(1, getCounterValue("reindexing.jobs", "status", "started")); + assertEquals(1.0, getGaugeValue("reindexing.jobs.active")); + } + + @Test + @DisplayName("recordJobCompleted increments counter and decrements active gauge") + void testRecordJobCompleted() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + Timer.Sample sample = metrics.startJobTimer(); + metrics.recordJobCompleted(sample); + + assertEquals(1, getCounterValue("reindexing.jobs", "status", "completed")); + assertEquals(0.0, getGaugeValue("reindexing.jobs.active")); + + Timer timer = + meterRegistry.find("reindexing.job.duration").tag("status", "completed").timer(); + assertNotNull(timer); + assertEquals(1, timer.count()); + } + + @Test + @DisplayName("recordJobFailed increments counter and decrements active gauge") + void testRecordJobFailed() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + Timer.Sample sample = metrics.startJobTimer(); + metrics.recordJobFailed(sample); + + assertEquals(1, getCounterValue("reindexing.jobs", "status", "failed")); + assertEquals(0.0, getGaugeValue("reindexing.jobs.active")); + + Timer timer = meterRegistry.find("reindexing.job.duration").tag("status", "failed").timer(); + assertNotNull(timer); + assertEquals(1, timer.count()); + } + + @Test + @DisplayName("recordJobStopped increments counter and decrements active gauge") + void testRecordJobStopped() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + Timer.Sample sample = metrics.startJobTimer(); + metrics.recordJobStopped(sample); + + assertEquals(1, getCounterValue("reindexing.jobs", "status", "stopped")); + assertEquals(0.0, getGaugeValue("reindexing.jobs.active")); + + Timer timer = meterRegistry.find("reindexing.job.duration").tag("status", "stopped").timer(); + assertNotNull(timer); + assertEquals(1, timer.count()); + } + + @Test + @DisplayName("recordJobCompleted handles null sample gracefully") + void testRecordJobCompletedNullSample() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + metrics.recordJobCompleted(null); + + assertEquals(1, getCounterValue("reindexing.jobs", "status", "completed")); + Timer timer = + meterRegistry.find("reindexing.job.duration").tag("status", "completed").timer(); + assertNotNull(timer); + assertEquals(0, timer.count()); + } + + @Test + @DisplayName("Multiple jobs tracked correctly") + void testMultipleJobs() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordJobStarted(); + metrics.recordJobStarted(); + assertEquals(2.0, getGaugeValue("reindexing.jobs.active")); + + metrics.recordJobCompleted(null); + assertEquals(1.0, getGaugeValue("reindexing.jobs.active")); + + metrics.recordJobFailed(null); + assertEquals(0.0, getGaugeValue("reindexing.jobs.active")); + + assertEquals(2, getCounterValue("reindexing.jobs", "status", "started")); + assertEquals(1, getCounterValue("reindexing.jobs", "status", "completed")); + assertEquals(1, getCounterValue("reindexing.jobs", "status", "failed")); + } + } + + @Nested + @DisplayName("Stage Counters") + class StageCounterTests { + + @Test + @DisplayName("recordStageSuccess creates counter with correct tags") + void testRecordStageSuccess() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordStageSuccess("reader", "table", 10); + + Counter counter = + meterRegistry + .find("reindexing.stage.success") + .tag("stage", "reader") + .tag("entity_type", "table") + .counter(); + assertNotNull(counter); + assertEquals(10.0, counter.count()); + } + + @Test + @DisplayName("recordStageFailed creates counter with correct tags") + void testRecordStageFailed() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordStageFailed("sink", "dashboard", 3); + + Counter counter = + meterRegistry + .find("reindexing.stage.failed") + .tag("stage", "sink") + .tag("entity_type", "dashboard") + .counter(); + assertNotNull(counter); + assertEquals(3.0, counter.count()); + } + + @Test + @DisplayName("recordStageWarnings creates counter with correct tags") + void testRecordStageWarnings() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordStageWarnings("reader", "pipeline", 5); + + Counter counter = + meterRegistry + .find("reindexing.stage.warnings") + .tag("stage", "reader") + .tag("entity_type", "pipeline") + .counter(); + assertNotNull(counter); + assertEquals(5.0, counter.count()); + } + + @Test + @DisplayName("Stage counters accumulate across calls") + void testStageCountersAccumulate() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordStageSuccess("process", "table", 100); + metrics.recordStageSuccess("process", "table", 50); + + Counter counter = + meterRegistry + .find("reindexing.stage.success") + .tag("stage", "process") + .tag("entity_type", "table") + .counter(); + assertNotNull(counter); + assertEquals(150.0, counter.count()); + } + + @Test + @DisplayName("Different entity types produce separate counters") + void testDifferentEntityTypes() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordStageSuccess("reader", "table", 10); + metrics.recordStageSuccess("reader", "topic", 20); + + Counter tableCounter = + meterRegistry + .find("reindexing.stage.success") + .tag("stage", "reader") + .tag("entity_type", "table") + .counter(); + Counter topicCounter = + meterRegistry + .find("reindexing.stage.success") + .tag("stage", "reader") + .tag("entity_type", "topic") + .counter(); + assertNotNull(tableCounter); + assertNotNull(topicCounter); + assertEquals(10.0, tableCounter.count()); + assertEquals(20.0, topicCounter.count()); + } + } + + @Nested + @DisplayName("Bulk Request Metrics") + class BulkRequestTests { + + @Test + @DisplayName("Bulk request timer records successful request") + void testBulkRequestSuccess() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + Timer.Sample sample = metrics.startBulkRequestTimer(); + metrics.recordBulkRequestCompleted(sample, true); + + Timer timer = meterRegistry.find("reindexing.bulk.duration").tag("success", "true").timer(); + assertNotNull(timer); + assertEquals(1, timer.count()); + } + + @Test + @DisplayName("Bulk request timer records failed request") + void testBulkRequestFailure() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + Timer.Sample sample = metrics.startBulkRequestTimer(); + metrics.recordBulkRequestCompleted(sample, false); + + Timer timer = meterRegistry.find("reindexing.bulk.duration").tag("success", "false").timer(); + assertNotNull(timer); + assertEquals(1, timer.count()); + } + + @Test + @DisplayName("recordBulkRequestCompleted handles null sample gracefully") + void testBulkRequestNullSample() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordBulkRequestCompleted(null, true); + + Timer timer = meterRegistry.find("reindexing.bulk.duration").tag("success", "true").timer(); + assertNotNull(timer); + assertEquals(0, timer.count()); + } + + @Test + @DisplayName("Payload size is recorded in distribution summary") + void testRecordPayloadSize() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordPayloadSize(1024); + metrics.recordPayloadSize(2048); + + DistributionSummary summary = meterRegistry.find("reindexing.bulk.payload.size").summary(); + assertNotNull(summary); + assertEquals(2, summary.count()); + assertEquals(3072.0, summary.totalAmount()); + } + + @Test + @DisplayName("Pending bulk requests gauge tracks increment and decrement") + void testPendingBulkRequests() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.incrementPendingBulkRequests(); + metrics.incrementPendingBulkRequests(); + assertEquals(2.0, getGaugeValue("reindexing.sink.pending")); + + metrics.decrementPendingBulkRequests(); + assertEquals(1.0, getGaugeValue("reindexing.sink.pending")); + } + } + + @Nested + @DisplayName("Backpressure") + class BackpressureTests { + + @Test + @DisplayName("recordBackpressureEvent increments counter") + void testRecordBackpressureEvent() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordBackpressureEvent(); + metrics.recordBackpressureEvent(); + metrics.recordBackpressureEvent(); + + assertEquals(3, getCounterValue("reindexing.backpressure.events")); + } + } + + @Nested + @DisplayName("Promotion Metrics") + class PromotionTests { + + @Test + @DisplayName("recordPromotionSuccess creates counter with correct tags") + void testRecordPromotionSuccess() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordPromotionSuccess("table"); + + Counter counter = + meterRegistry + .find("reindexing.promotion") + .tag("entity_type", "table") + .tag("result", "success") + .counter(); + assertNotNull(counter); + assertEquals(1.0, counter.count()); + } + + @Test + @DisplayName("recordPromotionFailure creates counter with correct tags") + void testRecordPromotionFailure() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordPromotionFailure("dashboard"); + + Counter counter = + meterRegistry + .find("reindexing.promotion") + .tag("entity_type", "dashboard") + .tag("result", "failure") + .counter(); + assertNotNull(counter); + assertEquals(1.0, counter.count()); + } + + @Test + @DisplayName("Promotion success and failure tracked independently per entity type") + void testPromotionPerEntityType() { + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + metrics.recordPromotionSuccess("table"); + metrics.recordPromotionSuccess("table"); + metrics.recordPromotionFailure("table"); + metrics.recordPromotionSuccess("topic"); + + Counter tableSuccess = + meterRegistry + .find("reindexing.promotion") + .tag("entity_type", "table") + .tag("result", "success") + .counter(); + Counter tableFailure = + meterRegistry + .find("reindexing.promotion") + .tag("entity_type", "table") + .tag("result", "failure") + .counter(); + Counter topicSuccess = + meterRegistry + .find("reindexing.promotion") + .tag("entity_type", "topic") + .tag("result", "success") + .counter(); + + assertNotNull(tableSuccess); + assertNotNull(tableFailure); + assertNotNull(topicSuccess); + assertEquals(2.0, tableSuccess.count()); + assertEquals(1.0, tableFailure.count()); + assertEquals(1.0, topicSuccess.count()); + } + } + + private long getCounterValue(String name) { + Counter counter = meterRegistry.find(name).counter(); + return counter != null ? (long) counter.count() : 0; + } + + private long getCounterValue(String name, String tagKey, String tagValue) { + Counter counter = meterRegistry.find(name).tag(tagKey, tagValue).counter(); + return counter != null ? (long) counter.count() : 0; + } + + private double getGaugeValue(String name) { + Gauge gauge = meterRegistry.find(name).gauge(); + return gauge != null ? gauge.value() : 0.0; + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppTest.java index 3047c383061..461eff3092d 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppTest.java @@ -47,8 +47,6 @@ import org.openmetadata.schema.api.services.CreateMessagingService; import org.openmetadata.schema.api.services.CreateMessagingService.MessagingServiceType; import org.openmetadata.schema.entity.app.App; import org.openmetadata.schema.entity.app.AppRunRecord; -import org.openmetadata.schema.entity.app.FailureContext; -import org.openmetadata.schema.entity.app.SuccessContext; import org.openmetadata.schema.entity.data.Database; import org.openmetadata.schema.entity.data.DatabaseSchema; import org.openmetadata.schema.entity.data.Table; @@ -57,9 +55,7 @@ import org.openmetadata.schema.entity.services.DatabaseService; import org.openmetadata.schema.entity.services.MessagingService; import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration; import org.openmetadata.schema.system.EventPublisherJob; -import org.openmetadata.schema.system.IndexingError; import org.openmetadata.schema.system.Stats; -import org.openmetadata.schema.system.StepStats; import org.openmetadata.schema.type.AccessDetails; import org.openmetadata.schema.type.Column; import org.openmetadata.schema.type.ColumnDataType; @@ -794,96 +790,42 @@ class SearchIndexAppTest extends OpenMetadataApplicationTest { } @Test - void testJobCompletionStatus() throws Exception { + void testOrchestratorSendUpdatesPopulatesAppRunRecord() { try (MockedStatic wsMock = mockStatic(WebSocketManager.class)) { wsMock.when(WebSocketManager::getInstance).thenReturn(webSocketManager); - searchIndexApp.init( - new App() - .withName("SearchIndexingApplication") - .withAppConfiguration(JsonUtils.convertValue(testJobData, Object.class))); + UUID appId = UUID.randomUUID(); + AppRunRecord record = new AppRunRecord(); + record.setStatus(AppRunRecord.Status.RUNNING); - EventPublisherJob jobData = searchIndexApp.getJobData(); - jobData.setStatus(EventPublisherJob.Status.RUNNING); + OrchestratorContext orchCtx = mock(OrchestratorContext.class); + when(orchCtx.getJobRecord()).thenReturn(record); + when(orchCtx.getJobName()).thenReturn("TestJob"); + when(orchCtx.getAppConfigJson()).thenReturn(JsonUtils.pojoToJson(testJobData)); + when(orchCtx.getAppId()).thenReturn(appId); + when(orchCtx.createProgressListener(any())) + .thenReturn(mock(ReindexingProgressListener.class)); + when(orchCtx.createReindexingContext(false)).thenReturn(mock(ReindexingJobContext.class)); - var method = - SearchIndexApp.class.getDeclaredMethod( - "sendUpdates", JobExecutionContext.class, boolean.class); - method.setAccessible(true); + CollectionDAO.SearchIndexFailureDAO failureDAO = + mock(CollectionDAO.SearchIndexFailureDAO.class); + when(collectionDAO.searchIndexFailureDAO()).thenReturn(failureDAO); + when(failureDAO.deleteAll()).thenReturn(0); + when(failureDAO.countByJobId(anyString())).thenReturn(0); - if (jobData.getStatus() == EventPublisherJob.Status.RUNNING) { - jobData.setStatus(EventPublisherJob.Status.COMPLETED); - method.invoke(searchIndexApp, jobExecutionContext, true); - } + ReindexingOrchestrator orch = + new ReindexingOrchestrator(collectionDAO, searchRepository, orchCtx); - assertEquals(EventPublisherJob.Status.COMPLETED, jobData.getStatus()); - } - } + EventPublisherJob emptyJob = + new EventPublisherJob() + .withEntities(Set.of()) + .withBatchSize(100) + .withRecreateIndex(false); - @Test - void testWebSocketThrottling() throws Exception { - try (MockedStatic wsMock = mockStatic(WebSocketManager.class)) { - wsMock.when(WebSocketManager::getInstance).thenReturn(webSocketManager); + orch.run(emptyJob); - searchIndexApp.init( - new App() - .withName("SearchIndexingApplication") - .withAppConfiguration(JsonUtils.convertValue(testJobData, Object.class))); - - var method = - SearchIndexApp.class.getDeclaredMethod( - "sendUpdates", JobExecutionContext.class, boolean.class); - method.setAccessible(true); - - method.invoke(searchIndexApp, jobExecutionContext, false); - method.invoke(searchIndexApp, jobExecutionContext, false); - method.invoke(searchIndexApp, jobExecutionContext, false); - method.invoke(searchIndexApp, jobExecutionContext, true); - } - } - - @Test - void testAppRunRecordCreation() { - try (MockedStatic wsMock = mockStatic(WebSocketManager.class)) { - wsMock.when(WebSocketManager::getInstance).thenReturn(webSocketManager); - - searchIndexApp.init( - new App() - .withName("SearchIndexingApplication") - .withAppConfiguration(JsonUtils.convertValue(testJobData, Object.class))); - - EventPublisherJob jobData = searchIndexApp.getJobData(); - - IndexingError error = - new IndexingError() - .withErrorSource(IndexingError.ErrorSource.SINK) - .withMessage("Test error") - .withFailedCount(5); - jobData.setFailure(error); - jobData.setStatus(EventPublisherJob.Status.ACTIVE_ERROR); - - Stats stats = - new Stats().withJobStats(new StepStats().withSuccessRecords(95).withFailedRecords(5)); - jobData.setStats(stats); - - AppRunRecord mockRecord = mock(AppRunRecord.class); - lenient().when(mockRecord.getStatus()).thenReturn(AppRunRecord.Status.FAILED); - lenient().when(mockRecord.getFailureContext()).thenReturn(new FailureContext()); - lenient().when(mockRecord.getSuccessContext()).thenReturn(new SuccessContext()); - - try { - var method = - SearchIndexApp.class.getDeclaredMethod( - "updateRecordToDbAndNotify", JobExecutionContext.class); - method.setAccessible(true); - method.invoke(searchIndexApp, jobExecutionContext); - assertEquals(IndexingError.ErrorSource.SINK, jobData.getFailure().getErrorSource()); - assertEquals("Test error", jobData.getFailure().getMessage()); - assertEquals(5, jobData.getFailure().getFailedCount()); - - } catch (Exception e) { - LOG.debug("Expected exception during partial mocking: {}", e.getMessage()); - } + assertEquals(EventPublisherJob.Status.COMPLETED, orch.getJobData().getStatus()); + assertEquals(AppRunRecord.Status.COMPLETED.value(), record.getStatus().value()); } } @@ -1319,14 +1261,9 @@ class SearchIndexAppTest extends OpenMetadataApplicationTest { @Test void testInitializeTotalRecords() { - App testApp = - new App() - .withName("SearchIndexingApplication") - .withAppConfiguration(JsonUtils.convertValue(testJobData, Object.class)); + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); - searchIndexApp.init(testApp); - - Stats stats = searchIndexApp.initializeTotalRecords(Set.of("table", "user")); + Stats stats = executor.initializeTotalRecords(Set.of("table", "user")); assertNotNull(stats); assertNotNull(stats.getJobStats()); assertNotNull(stats.getReaderStats()); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureIntegrationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureIntegrationTest.java index 073542745df..3111186c46f 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureIntegrationTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureIntegrationTest.java @@ -66,13 +66,33 @@ class SearchIndexFailureIntegrationTest { // Simulate what happens when BulkSink fails to index entities BulkSink.FailureCallback callback = - (entityType, entityId, entityFqn, errorMessage) -> + (entityType, entityId, entityFqn, errorMessage, stage) -> { + if (stage == IndexingFailureRecorder.FailureStage.PROCESS) { + recorder.recordProcessFailure(entityType, entityId, entityFqn, errorMessage); + } else { recorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + } + }; // Simulate 3 sink failures - callback.onFailure("table", "uuid-1", "db.schema.table1", "Mapping error"); - callback.onFailure("table", "uuid-2", "db.schema.table2", "Document too large"); - callback.onFailure("dashboard", "uuid-3", "service.dashboard1", "Index not found"); + callback.onFailure( + "table", + "uuid-1", + "db.schema.table1", + "Mapping error", + IndexingFailureRecorder.FailureStage.SINK); + callback.onFailure( + "table", + "uuid-2", + "db.schema.table2", + "Document too large", + IndexingFailureRecorder.FailureStage.SINK); + callback.onFailure( + "dashboard", + "uuid-3", + "service.dashboard1", + "Index not found", + IndexingFailureRecorder.FailureStage.SINK); // Flush to capture recorder.flush(); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/StatsThreadSafetyTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/StatsThreadSafetyTest.java new file mode 100644 index 00000000000..7ea8431e000 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/StatsThreadSafetyTest.java @@ -0,0 +1,196 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.system.EntityStats; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; + +@DisplayName("Stats Thread Safety Tests") +class StatsThreadSafetyTest { + + private Stats createTestStats(Set entityTypes) { + Stats stats = new Stats(); + stats.setEntityStats(new EntityStats()); + stats.setJobStats(new StepStats()); + stats.setReaderStats(new StepStats()); + stats.setSinkStats(new StepStats()); + + int total = 0; + for (String entityType : entityTypes) { + int entityTotal = 1000; + total += entityTotal; + StepStats es = new StepStats(); + es.setTotalRecords(entityTotal); + es.setSuccessRecords(0); + es.setFailedRecords(0); + stats.getEntityStats().getAdditionalProperties().put(entityType, es); + } + stats.getJobStats().setTotalRecords(total); + stats.getJobStats().setSuccessRecords(0); + stats.getJobStats().setFailedRecords(0); + stats.getReaderStats().setTotalRecords(total); + stats.getReaderStats().setSuccessRecords(0); + stats.getReaderStats().setFailedRecords(0); + stats.getReaderStats().setWarningRecords(0); + stats.getSinkStats().setTotalRecords(0); + stats.getSinkStats().setSuccessRecords(0); + stats.getSinkStats().setFailedRecords(0); + return stats; + } + + @Test + @DisplayName("concurrent updateStats calls produce consistent entity and job stats") + void concurrentUpdateStats() throws InterruptedException { + Set entities = Set.of("table", "user"); + AtomicReference statsRef = new AtomicReference<>(createTestStats(entities)); + Stats stats = statsRef.get(); + + int threadCount = 8; + int updatesPerThread = 100; + CountDownLatch latch = new CountDownLatch(threadCount); + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + + for (int t = 0; t < threadCount; t++) { + final String entityType = (t % 2 == 0) ? "table" : "user"; + executor.submit( + () -> { + try { + for (int i = 0; i < updatesPerThread; i++) { + synchronized (statsRef) { + StepStats es = stats.getEntityStats().getAdditionalProperties().get(entityType); + if (es != null) { + es.setSuccessRecords(es.getSuccessRecords() + 1); + } + int totalSuccess = + stats.getEntityStats().getAdditionalProperties().values().stream() + .mapToInt(StepStats::getSuccessRecords) + .sum(); + stats.getJobStats().setSuccessRecords(totalSuccess); + } + } + } finally { + latch.countDown(); + } + }); + } + + latch.await(10, TimeUnit.SECONDS); + executor.shutdown(); + + assertNotNull(stats.getJobStats()); + int tableSuccess = + stats.getEntityStats().getAdditionalProperties().get("table").getSuccessRecords(); + int userSuccess = + stats.getEntityStats().getAdditionalProperties().get("user").getSuccessRecords(); + int jobSuccess = stats.getJobStats().getSuccessRecords(); + + assertEquals(threadCount * updatesPerThread, tableSuccess + userSuccess); + assertEquals(tableSuccess + userSuccess, jobSuccess); + } + + @Test + @DisplayName("concurrent reader stats updates do not lose updates") + void concurrentReaderStats() throws InterruptedException { + Stats stats = createTestStats(Set.of("table")); + + int threadCount = 8; + int updatesPerThread = 100; + CountDownLatch latch = new CountDownLatch(threadCount); + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + + for (int t = 0; t < threadCount; t++) { + executor.submit( + () -> { + try { + for (int i = 0; i < updatesPerThread; i++) { + synchronized (stats) { + StepStats rs = stats.getReaderStats(); + rs.setSuccessRecords( + (rs.getSuccessRecords() != null ? rs.getSuccessRecords() : 0) + 1); + rs.setFailedRecords( + (rs.getFailedRecords() != null ? rs.getFailedRecords() : 0) + 0); + rs.setWarningRecords( + (rs.getWarningRecords() != null ? rs.getWarningRecords() : 0) + 0); + } + } + } finally { + latch.countDown(); + } + }); + } + + latch.await(10, TimeUnit.SECONDS); + executor.shutdown(); + + int expectedSuccess = threadCount * updatesPerThread; + assertEquals(expectedSuccess, stats.getReaderStats().getSuccessRecords()); + assertEquals(0, stats.getReaderStats().getFailedRecords()); + } + + @Test + @DisplayName("reconciler invariant: job total >= job success + job failed") + void reconcilerInvariant() throws InterruptedException { + Stats stats = createTestStats(Set.of("table", "user")); + + int threadCount = 4; + int updatesPerThread = 200; + CountDownLatch latch = new CountDownLatch(threadCount); + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + + for (int t = 0; t < threadCount; t++) { + final int tid = t; + executor.submit( + () -> { + try { + for (int i = 0; i < updatesPerThread; i++) { + String entityType = (tid % 2 == 0) ? "table" : "user"; + synchronized (stats) { + StepStats es = stats.getEntityStats().getAdditionalProperties().get(entityType); + if (es != null) { + if (i % 10 == 0) { + es.setFailedRecords(es.getFailedRecords() + 1); + } else { + es.setSuccessRecords(es.getSuccessRecords() + 1); + } + } + int totalSuccess = + stats.getEntityStats().getAdditionalProperties().values().stream() + .mapToInt(StepStats::getSuccessRecords) + .sum(); + int totalFailed = + stats.getEntityStats().getAdditionalProperties().values().stream() + .mapToInt(StepStats::getFailedRecords) + .sum(); + stats.getJobStats().setSuccessRecords(totalSuccess); + stats.getJobStats().setFailedRecords(totalFailed); + } + } + } finally { + latch.countDown(); + } + }); + } + + latch.await(10, TimeUnit.SECONDS); + executor.shutdown(); + + StepStats jobStats = stats.getJobStats(); + int success = jobStats.getSuccessRecords(); + int failed = jobStats.getFailedRecords(); + int total = jobStats.getTotalRecords(); + + assertTrue(total >= success + failed, "Total must be >= success + failed"); + assertEquals(threadCount * updatesPerThread, success + failed); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinatorTest.java index 5aecdc3daa7..ef06a1419eb 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinatorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinatorTest.java @@ -104,7 +104,7 @@ class DistributedSearchIndexCoordinatorTest { Set entities = Set.of("table", "user"); EventPublisherJob jobConfig = new EventPublisherJob().withEntities(entities).withBatchSize(500); - when(partitionCalculator.getEntityCounts(entities)) + when(partitionCalculator.getEntityCounts(entities, null)) .thenReturn(java.util.Map.of("table", 10000L, "user", 5000L)); SearchIndexJob job = coordinator.createJob(entities, jobConfig, "admin"); @@ -194,7 +194,7 @@ class DistributedSearchIndexCoordinatorTest { .cursor(5000) .build()); - when(partitionCalculator.calculatePartitions(jobId, entities)).thenReturn(mockPartitions); + when(partitionCalculator.calculatePartitions(jobId, entities, null)).thenReturn(mockPartitions); SearchIndexJob result = coordinator.initializePartitions(jobId); @@ -275,7 +275,8 @@ class DistributedSearchIndexCoordinatorTest { // Atomic claim succeeds when(partitionDAO.claimNextPartitionAtomic(eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn(1); - when(partitionDAO.findLatestClaimedPartition(jobId.toString(), TEST_SERVER_ID)) + when(partitionDAO.findLatestClaimedPartition( + eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn( new SearchIndexPartitionRecord( partitionId.toString(), @@ -326,7 +327,7 @@ class DistributedSearchIndexCoordinatorTest { Optional result = coordinator.claimNextPartition(jobId); assertFalse(result.isPresent()); - verify(partitionDAO, never()).findLatestClaimedPartition(anyString(), anyString()); + verify(partitionDAO, never()).findLatestClaimedPartition(anyString(), anyString(), anyLong()); } @Test @@ -354,7 +355,8 @@ class DistributedSearchIndexCoordinatorTest { when(partitionDAO.claimNextPartitionAtomic(eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn(1); - when(partitionDAO.findLatestClaimedPartition(jobId.toString(), TEST_SERVER_ID)) + when(partitionDAO.findLatestClaimedPartition( + eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn( new SearchIndexPartitionRecord( partitionId.toString(), @@ -405,7 +407,7 @@ class DistributedSearchIndexCoordinatorTest { // Should return empty - lost the race assertFalse(result.isPresent()); // Should NOT call findLatestClaimedPartition since claim failed - verify(partitionDAO, never()).findLatestClaimedPartition(anyString(), anyString()); + verify(partitionDAO, never()).findLatestClaimedPartition(anyString(), anyString(), anyLong()); } @Test @@ -418,7 +420,8 @@ class DistributedSearchIndexCoordinatorTest { when(partitionDAO.claimNextPartitionAtomic(eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn(1); - when(partitionDAO.findLatestClaimedPartition(jobId.toString(), TEST_SERVER_ID)) + when(partitionDAO.findLatestClaimedPartition( + eq(jobId.toString()), eq(TEST_SERVER_ID), anyLong())) .thenReturn( new SearchIndexPartitionRecord( partitionId.toString(), diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexIntegrationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexIntegrationTest.java index 61864f35794..f9300ceca56 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexIntegrationTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexIntegrationTest.java @@ -249,7 +249,7 @@ class DistributedSearchIndexIntegrationTest extends OpenMetadataApplicationTest assertEquals(1, claimed, "Should claim exactly one partition"); SearchIndexPartitionRecord claimedPartition = - partitionDAO.findLatestClaimedPartition(jobId, "server-1"); + partitionDAO.findLatestClaimedPartition(jobId, "server-1", now); assertNotNull(claimedPartition, "Should find the claimed partition"); assertEquals("PROCESSING", claimedPartition.status()); assertEquals("server-1", claimedPartition.assignedServer()); @@ -778,7 +778,8 @@ class DistributedSearchIndexIntegrationTest extends OpenMetadataApplicationTest int processedPartitions = 0; while (partitionDAO.claimNextPartitionAtomic(jobId, serverId, now) > 0) { - SearchIndexPartitionRecord claimed = partitionDAO.findLatestClaimedPartition(jobId, serverId); + SearchIndexPartitionRecord claimed = + partitionDAO.findLatestClaimedPartition(jobId, serverId, now); assertNotNull(claimed, "Should have a claimed partition"); partitionDAO.updateProgress( diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json index 25a5969c5ca..4f7ca974f75 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json @@ -109,6 +109,21 @@ "default": 10000, "minimum": 1000, "maximum": 50000 + }, + "timeSeriesMaxDays": { + "title": "Time Series Max Days", + "description": "Maximum age in days for time series data during reindexing. Default 0 (index all data). Set to a positive value like 15 to limit to recent data only.", + "type": "integer", + "default": 0, + "minimum": -1 + }, + "timeSeriesEntityDays": { + "title": "Time Series Entity Days Override", + "description": "Per-entity-type override for time series max days. Keys are entity type names (e.g. testCaseResult, queryCostRecord), values are number of days. Entities not listed here use the default Time Series Max Days value.", + "type": "object", + "additionalProperties": { + "type": "integer" + } } }, "additionalProperties": false diff --git a/openmetadata-spec/src/main/resources/json/schema/system/eventPublisherJob.json b/openmetadata-spec/src/main/resources/json/schema/system/eventPublisherJob.json index 6d277003656..a9986ab8e41 100644 --- a/openmetadata-spec/src/main/resources/json/schema/system/eventPublisherJob.json +++ b/openmetadata-spec/src/main/resources/json/schema/system/eventPublisherJob.json @@ -232,6 +232,22 @@ "default": 10000, "minimum": 1000, "maximum": 50000 + }, + "timeSeriesMaxDays": { + "title": "Time Series Max Days", + "description": "Maximum age in days for time series data during reindexing. Only records from the last N days will be indexed. Default 0 (index all data). Set to a positive value like 15 to limit to recent data.", + "type": "integer", + "default": 0, + "minimum": -1 + }, + "timeSeriesEntityDays": { + "title": "Time Series Entity Days Override", + "description": "Per-entity-type override for time series max days. Keys are entity type names, values are number of days. Entities not in this map use timeSeriesMaxDays as default.", + "type": "object", + "existingJavaType": "java.util.Map", + "additionalProperties": { + "type": "integer" + } } }, "additionalProperties": false diff --git a/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md b/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md index d099fb672b7..bae3f5a4c6f 100644 --- a/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md +++ b/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md @@ -101,4 +101,18 @@ $$section Number of entities per partition for distributed indexing. Smaller values create more partitions for better distribution across servers. Range: 1000-50000. +$$ + +$$section +### Time Series Max Days $(id="timeSeriesMaxDays") + +Maximum age in days for time series data during reindexing. Default 0 (index all data). Set to a positive value like 15 to limit to recent data only. + +$$ + +$$section +### Time Series Entity Days Override $(id="timeSeriesEntityDays") + +Per-entity-type override for time series max days. Keys are entity type names (e.g. testCaseResult, queryCostRecord), values are number of days. Entities not listed here use the default Time Series Max Days value. + $$ \ No newline at end of file diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/api/services/ingestionPipelines/createIngestionPipeline.ts b/openmetadata-ui/src/main/resources/ui/src/generated/api/services/ingestionPipelines/createIngestionPipeline.ts index 1d79ffb75e8..eccc3fcc9d2 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/api/services/ingestionPipelines/createIngestionPipeline.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/api/services/ingestionPipelines/createIngestionPipeline.ts @@ -1090,6 +1090,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/app.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/app.ts index 07e1ef5e956..8531b0d4eae 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/app.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/app.ts @@ -355,6 +355,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/configuration/internal/searchIndexingAppConfig.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/configuration/internal/searchIndexingAppConfig.ts index 5185a746f47..646dc346488 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/configuration/internal/searchIndexingAppConfig.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/configuration/internal/searchIndexingAppConfig.ts @@ -72,6 +72,17 @@ export interface SearchIndexingAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Application Type */ diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/appMarketPlaceDefinition.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/appMarketPlaceDefinition.ts index 1f0d4eaa4f3..4cf6bb78b25 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/appMarketPlaceDefinition.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/appMarketPlaceDefinition.ts @@ -336,6 +336,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/createAppMarketPlaceDefinitionReq.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/createAppMarketPlaceDefinitionReq.ts index e79ea2138e4..ee42ee807ae 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/createAppMarketPlaceDefinitionReq.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/applications/marketplace/createAppMarketPlaceDefinitionReq.ts @@ -293,6 +293,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/services/ingestionPipelines/ingestionPipeline.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/services/ingestionPipelines/ingestionPipeline.ts index 55c815d58be..c3f3d1cca0c 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/services/ingestionPipelines/ingestionPipeline.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/services/ingestionPipelines/ingestionPipeline.ts @@ -1627,6 +1627,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/application.ts b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/application.ts index a4316138cdf..d7706965af9 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/application.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/application.ts @@ -196,6 +196,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/applicationPipeline.ts b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/applicationPipeline.ts index 204a22352e9..03be0849568 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/applicationPipeline.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/applicationPipeline.ts @@ -181,6 +181,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/workflow.ts b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/workflow.ts index da8985394c5..b418404ee92 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/workflow.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/metadataIngestion/workflow.ts @@ -5545,6 +5545,17 @@ export interface CollateAIAppConfig { * Recreate Indexes with updated Language */ searchIndexMappingLanguage?: SearchIndexMappingLanguage; + /** + * Per-entity-type override for time series max days. Keys are entity type names (e.g. + * testCaseResult, queryCostRecord), values are number of days. Entities not listed here use + * the default Time Series Max Days value. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Default 0 (index all data). + * Set to a positive value like 15 to limit to recent data only. + */ + timeSeriesMaxDays?: number; /** * Enable distributed indexing to scale reindexing across multiple servers with fault * tolerance and parallel processing diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/system/eventPublisherJob.ts b/openmetadata-ui/src/main/resources/ui/src/generated/system/eventPublisherJob.ts index b2dfb063e5b..a9d8ca9fb55 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/system/eventPublisherJob.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/system/eventPublisherJob.ts @@ -101,8 +101,19 @@ export interface EventPublisherJob { /** * This schema publisher run job status. */ - status?: Status; - timestamp?: number; + status?: Status; + /** + * Per-entity-type override for time series max days. Keys are entity type names, values are + * number of days. Entities not in this map use timeSeriesMaxDays as default. + */ + timeSeriesEntityDays?: { [key: string]: number }; + /** + * Maximum age in days for time series data during reindexing. Only records from the last N + * days will be indexed. Default 0 (index all data). Set to a positive value like 15 to + * limit to recent data. + */ + timeSeriesMaxDays?: number; + timestamp?: number; /** * Enable distributed indexing across multiple servers. When enabled, reindexing work is * partitioned and can be processed by multiple servers concurrently with crash recovery diff --git a/openmetadata-ui/src/main/resources/ui/src/utils/ApplicationSchemas/SearchIndexingApplication.json b/openmetadata-ui/src/main/resources/ui/src/utils/ApplicationSchemas/SearchIndexingApplication.json index 6a7fb62d446..6531eba4c8f 100644 --- a/openmetadata-ui/src/main/resources/ui/src/utils/ApplicationSchemas/SearchIndexingApplication.json +++ b/openmetadata-ui/src/main/resources/ui/src/utils/ApplicationSchemas/SearchIndexingApplication.json @@ -161,6 +161,21 @@ "default": 10000, "minimum": 1000, "maximum": 50000 + }, + "timeSeriesMaxDays": { + "title": "Time Series Max Days", + "description": "Maximum age in days for time series data during reindexing. Default 0 (index all data). Set to a positive value like 15 to limit to recent data only.", + "type": "integer", + "default": 0, + "minimum": -1 + }, + "timeSeriesEntityDays": { + "title": "Time Series Entity Days Override", + "description": "Per-entity-type override for time series max days. Keys are entity type names (e.g. testCaseResult, queryCostRecord), values are number of days. Entities not listed here use the default Time Series Max Days value.", + "type": "object", + "additionalProperties": { + "type": "integer" + } } }, "additionalProperties": false