forked from slackhq/astra
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
16d8f56
commit 3563980
Showing
17 changed files
with
259 additions
and
94 deletions.
There are no files selected for viewing
141 changes: 141 additions & 0 deletions
141
astra/src/main/java/com/slack/astra/bulkIngestApi/BulkLocalIngestApi.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,141 @@ | ||
| package com.slack.astra.bulkIngestApi; | ||
|
|
||
| import static com.linecorp.armeria.common.HttpStatus.CREATED; | ||
| import static com.linecorp.armeria.common.HttpStatus.INTERNAL_SERVER_ERROR; | ||
|
|
||
| import com.linecorp.armeria.common.HttpResponse; | ||
| import com.linecorp.armeria.server.annotation.Post; | ||
| import com.slack.astra.bulkIngestApi.opensearch.BulkApiRequestParser; | ||
| import com.slack.astra.chunkManager.ChunkManager; | ||
| import com.slack.astra.logstore.LogMessage; | ||
| import com.slack.astra.proto.schema.Schema; | ||
| import com.slack.service.murron.trace.Trace; | ||
| import java.nio.charset.StandardCharsets; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.concurrent.CompletableFuture; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| /** | ||
| * This class is responsible for defining the http endpoint behavior for the bulk ingest. It is | ||
| * expected to handle appropriate rate limiting, error handling, and submit the parsed messages to | ||
| * Kafka for ingestion. | ||
| */ | ||
| public class BulkLocalIngestApi { | ||
| private static final Logger LOG = LoggerFactory.getLogger(BulkLocalIngestApi.class); | ||
|
|
||
| // private final BulkIngestKafkaProducer bulkIngestKafkaProducer; | ||
| // private final DatasetRateLimitingService datasetRateLimitingService; | ||
| // private final MeterRegistry meterRegistry; | ||
| // private final Counter incomingByteTotal; | ||
| // private final Counter incomingDocsTotal; | ||
| // private final Timer bulkIngestTimer; | ||
| // private final String BULK_INGEST_INCOMING_BYTE_TOTAL = "astra_preprocessor_incoming_byte"; | ||
| // private final String BULK_INGEST_INCOMING_BYTE_DOCS = "astra_preprocessor_incoming_docs"; | ||
| // private final String BULK_INGEST_ERROR = "astra_preprocessor_error"; | ||
| // private final String BULK_INGEST_TIMER = "astra_preprocessor_bulk_ingest"; | ||
| // private final int rateLimitExceededErrorCode; | ||
| private final ChunkManager<LogMessage> chunkManager; | ||
| private final Schema.IngestSchema schema; | ||
|
|
||
| // private final Counter bulkIngestErrorCounter; | ||
|
|
||
| public BulkLocalIngestApi( | ||
| // MeterRegistry meterRegistry, | ||
| ChunkManager<LogMessage> chunkManager, Schema.IngestSchema schema) { | ||
|
|
||
| // this.bulkIngestKafkaProducer = bulkIngestKafkaProducer; | ||
| // this.datasetRateLimitingService = datasetRateLimitingService; | ||
| // this.meterRegistry = meterRegistry; | ||
| // this.incomingByteTotal = meterRegistry.counter(BULK_INGEST_INCOMING_BYTE_TOTAL); | ||
| // this.incomingDocsTotal = meterRegistry.counter(BULK_INGEST_INCOMING_BYTE_DOCS); | ||
| // this.bulkIngestTimer = meterRegistry.timer(BULK_INGEST_TIMER); | ||
| // if (rateLimitExceededErrorCode <= 0 || rateLimitExceededErrorCode > 599) { | ||
| // this.rateLimitExceededErrorCode = 400; | ||
| // } else { | ||
| // this.rateLimitExceededErrorCode = rateLimitExceededErrorCode; | ||
| // } | ||
| this.schema = schema; | ||
| this.chunkManager = chunkManager; | ||
| // this.bulkIngestErrorCounter = meterRegistry.counter(BULK_INGEST_ERROR); | ||
| } | ||
|
|
||
| @Post("/_local_bulk") | ||
| public HttpResponse addDocument(String bulkRequest) { | ||
| // 1. Astra does not support the concept of "updates". It's always an add. | ||
| // 2. The "index" is used as the span name | ||
| CompletableFuture<HttpResponse> future = new CompletableFuture<>(); | ||
| // Timer.Sample sample = Timer.start(meterRegistry); | ||
| // future.thenRun(() -> sample.stop(bulkIngestTimer)); | ||
|
|
||
| int count = 0; | ||
|
|
||
| try { | ||
| byte[] bulkRequestBytes = bulkRequest.getBytes(StandardCharsets.UTF_8); | ||
| // incomingByteTotal.increment(bulkRequestBytes.length); | ||
| Map<String, List<Trace.Span>> docs = Map.of(); | ||
| try { | ||
| docs = BulkApiRequestParser.parseRequest(bulkRequestBytes, schema); | ||
| } catch (Exception e) { | ||
| LOG.error("Request failed ", e); | ||
| // bulkIngestErrorCounter.increment(); | ||
| BulkIngestResponse response = new BulkIngestResponse(0, 0, e.getMessage()); | ||
| future.complete(HttpResponse.ofJson(INTERNAL_SERVER_ERROR, response)); | ||
| } | ||
| LOG.info("Parsed docs message: {}", docs); | ||
|
|
||
| // todo - our rate limiter doesn't have a way to acquire permits across multiple | ||
| // datasets | ||
| // so today as a limitation we reject any request that has documents against | ||
| // multiple indexes | ||
| // We think most indexing requests will be against 1 index | ||
| if (docs.keySet().size() > 1) { | ||
| BulkIngestResponse response = | ||
| new BulkIngestResponse(0, 0, "request must contain only 1 unique index"); | ||
| future.complete(HttpResponse.ofJson(INTERNAL_SERVER_ERROR, response)); | ||
| // bulkIngestErrorCounter.increment(); | ||
| return HttpResponse.of(future); | ||
| } | ||
|
|
||
| // for (Map.Entry<String, List<Trace.Span>> indexDocs : docs.entrySet()) { | ||
| // incomingDocsTotal.increment(indexDocs.getValue().size()); | ||
| // final String index = indexDocs.getKey(); | ||
| // if (!datasetRateLimitingService.tryAcquire(index, indexDocs.getValue())) { | ||
| // BulkIngestResponse response = new BulkIngestResponse(0, 0, "rate limit exceeded"); | ||
| // future.complete( | ||
| // HttpResponse.ofJson(HttpStatus.valueOf(rateLimitExceededErrorCode), | ||
| // response)); | ||
| // return HttpResponse.of(future); | ||
| // } | ||
| // } | ||
|
|
||
| // todo - explore the possibility of using the blocking task executor backed by virtual | ||
| // threads to fulfill this | ||
|
|
||
| for (Map.Entry<String, List<Trace.Span>> indexDocs : docs.entrySet()) { | ||
| for (Trace.Span span : indexDocs.getValue()) { | ||
| try { | ||
| chunkManager.addMessage(span, span.getSerializedSize(), String.valueOf(0), 12345, true); | ||
| count += 1; | ||
| // return HttpResponse.of(future); | ||
| } catch (Exception e) { | ||
| LOG.error("Request failed ", e); | ||
| // bulkIngestErrorCounter.increment(); | ||
| future.complete( | ||
| HttpResponse.ofJson( | ||
| INTERNAL_SERVER_ERROR, new BulkIngestResponse(0, 0, e.getMessage()))); | ||
| } | ||
| } | ||
| } | ||
| } catch (Exception e) { | ||
| LOG.error("Request failed ", e); | ||
| // bulkIngestErrorCounter.increment(); | ||
| BulkIngestResponse response = new BulkIngestResponse(0, 0, e.getMessage()); | ||
| future.complete(HttpResponse.ofJson(INTERNAL_SERVER_ERROR, response)); | ||
| } | ||
|
|
||
| future.complete(HttpResponse.ofJson(CREATED, new BulkIngestResponse(count, 0, ""))); | ||
| return HttpResponse.of(future); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.