diff --git a/distribution/pom.xml b/distribution/pom.xml index 6169d736af24..96b256e97220 100644 --- a/distribution/pom.xml +++ b/distribution/pom.xml @@ -139,6 +139,8 @@ -c org.apache.druid.extensions:druid-kafka-indexing-service -c + org.apache.druid.extensions:druid-kinesis-indexing-service + -c org.apache.druid.extensions:druid-lookups-cached-global -c org.apache.druid.extensions:druid-lookups-cached-single @@ -317,4 +319,4 @@ - + \ No newline at end of file diff --git a/docs/content/development/extensions-core/kafka-ingestion.md b/docs/content/development/extensions-core/kafka-ingestion.md index 2e0909706fae..aa771b7a27c1 100644 --- a/docs/content/development/extensions-core/kafka-ingestion.md +++ b/docs/content/development/extensions-core/kafka-ingestion.md @@ -28,7 +28,7 @@ The Kafka indexing service enables the configuration of *supervisors* on the Ove Kafka by managing the creation and lifetime of Kafka indexing tasks. These indexing tasks read events using Kafka's own partition and offset mechanism and are therefore able to provide guarantees of exactly-once ingestion. They are also able to read non-recent events from Kafka and are not subject to the window period considerations imposed on other -ingestion mechanisms. The supervisor oversees the state of the indexing tasks to coordinate handoffs, manage failures, +ingestion mechanisms using Tranquility. The supervisor oversees the state of the indexing tasks to coordinate handoffs, manage failures, and ensure that the scalability and replication requirements are maintained. This service is provided in the `druid-kafka-indexing-service` core extension (see @@ -125,8 +125,8 @@ A sample supervisor spec is shown below: |--------|-----------|---------| |`type`|The supervisor type, this should always be `kafka`.|yes| |`dataSchema`|The schema that will be used by the Kafka indexing task during ingestion, see [Ingestion Spec DataSchema](../../ingestion/ingestion-spec.html#dataschema).|yes| -|`tuningConfig`|A KafkaSupervisorTuningConfig to configure the supervisor and indexing tasks, see below.|no| |`ioConfig`|A KafkaSupervisorIOConfig to configure the supervisor and indexing tasks, see below.|yes| +|`tuningConfig`|A KafkaSupervisorTuningConfig to configure the supervisor and indexing tasks, see below.|no| ### KafkaSupervisorTuningConfig @@ -312,7 +312,7 @@ compatible because they have a different ingestion spec or partition allocation, supervisor will create a new set of tasks. In this way, the supervisors are persistent across overlord restarts and fail-overs. -A supervisor is stopped via the `POST /druid/indexer/v1/supervisor//shutdown` endpoint. This places a +A supervisor is stopped via the `POST /druid/indexer/v1/supervisor//terminate` endpoint. This places a tombstone marker in the database (to prevent the supervisor from being reloaded on a restart) and then gracefully shuts down the currently running supervisor. When a supervisor is shut down in this way, it will instruct its managed tasks to stop reading and begin publishing their segments immediately. The call to the shutdown endpoint will diff --git a/docs/content/development/extensions-core/kinesis-ingestion.md b/docs/content/development/extensions-core/kinesis-ingestion.md new file mode 100644 index 000000000000..8b3583b0d496 --- /dev/null +++ b/docs/content/development/extensions-core/kinesis-ingestion.md @@ -0,0 +1,392 @@ + + +--- +layout: doc_page +--- + +# Kinesis Indexing Service + +Similar to the [Kafka indexing service](./kafka-ingestion.html), the Kinesis indexing service enables the configuration of *supervisors* on the Overlord, which facilitate ingestion from +Kinesis by managing the creation and lifetime of Kinesis indexing tasks. These indexing tasks read events using Kinesis's own +Shards and Sequence Number mechanism and are therefore able to provide guarantees of exactly-once ingestion. They are also +able to read non-recent events from Kinesis and are not subject to the window period considerations imposed on other +ingestion mechanisms using Tranquility. The supervisor oversees the state of the indexing tasks to coordinate handoffs, manage failures, +and ensure that the scalability and replication requirements are maintained. + +The Kinesis indexing service is provided as the `druid-kinesis-indexing-service` core extension (see +[Including Extensions](../../operations/including-extensions.html)). Please note that this is +currently designated as an *experimental feature* and is subject to the usual +[experimental caveats](../experimental.html). + +## Submitting a Supervisor Spec + +The Kinesis indexing service requires that the `druid-kinesis-indexing-service` extension be loaded on both the overlord +and the middle managers. A supervisor for a dataSource is started by submitting a supervisor spec via HTTP POST to +`http://:/druid/indexer/v1/supervisor`, for example: + +``` +curl -X POST -H 'Content-Type: application/json' -d @supervisor-spec.json http://localhost:8090/druid/indexer/v1/supervisor +``` + +A sample supervisor spec is shown below: + +```json +{ + "type": "kinesis", + "dataSchema": { + "dataSource": "metrics-kinesis", + "parser": { + "type": "string", + "parseSpec": { + "format": "json", + "timestampSpec": { + "column": "timestamp", + "format": "auto" + }, + "dimensionsSpec": { + "dimensions": [], + "dimensionExclusions": [ + "timestamp", + "value" + ] + } + } + }, + "metricsSpec": [ + { + "name": "count", + "type": "count" + }, + { + "name": "value_sum", + "fieldName": "value", + "type": "doubleSum" + }, + { + "name": "value_min", + "fieldName": "value", + "type": "doubleMin" + }, + { + "name": "value_max", + "fieldName": "value", + "type": "doubleMax" + } + ], + "granularitySpec": { + "type": "uniform", + "segmentGranularity": "HOUR", + "queryGranularity": "NONE" + } + }, + "tuningConfig": { + "type": "kinesis", + "maxRowsPerSegment": 5000000 + }, + "ioConfig": { + "stream": "metrics", + "endpoint": "kinesis.us-east-1.amazonaws.com", + "taskCount": 1, + "replicas": 1, + "taskDuration": "PT1H", + "recordsPerFetch": 2000, + "fetchDelayMillis": 1000 + } +} +``` + +## Supervisor Configuration + +|Field|Description|Required| +|--------|-----------|---------| +|`type`|The supervisor type, this should always be `kinesis`.|yes| +|`dataSchema`|The schema that will be used by the Kinesis indexing task during ingestion, see [Ingestion Spec DataSchema](../../ingestion/ingestion-spec.html#dataschema).|yes| +|`ioConfig`|A KinesisSupervisorIOConfig to configure the supervisor and indexing tasks, see below.|yes| +|`tuningConfig`|A KinesisSupervisorTuningConfig to configure the supervisor and indexing tasks, see below.|no| + +### KinesisSupervisorTuningConfig + +The tuningConfig is optional and default parameters will be used if no tuningConfig is specified. + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|`type`|String|The indexing task type, this should always be `kinesis`.|yes| +|`maxRowsInMemory`|Integer|The number of rows to aggregate before persisting. This number is the post-aggregation rows, so it is not equivalent to the number of input events, but the number of aggregated rows that those events result in. This is used to manage the required JVM heap size. Maximum heap memory usage for indexing scales with maxRowsInMemory * (2 + maxPendingPersists).|no (default == 100000)| +|`maxBytesInMemory`|Long|The number of bytes to aggregate in heap memory before persisting. This is based on a rough estimate of memory usage and not actual usage. Normally this is computed internally and user does not need to set it. The maximum heap memory usage for indexing is maxBytesInMemory * (2 + maxPendingPersists). |no (default == One-sixth of max JVM memory)| +|`maxRowsPerSegment`|Integer|The number of rows to aggregate into a segment; this number is post-aggregation rows. Handoff will happen either if `maxRowsPerSegment` or `maxTotalRows` is hit or every `intermediateHandoffPeriod`, whichever happens earlier.|no (default == 5000000)| +|`maxTotalRows`|Long|The number of rows to aggregate across all segments; this number is post-aggregation rows. Handoff will happen either if `maxRowsPerSegment` or `maxTotalRows` is hit or every `intermediateHandoffPeriod`, whichever happens earlier.|no (default == unlimited)| +|`intermediatePersistPeriod`|ISO8601 Period|The period that determines the rate at which intermediate persists occur.|no (default == PT10M)| +|`maxPendingPersists`|Integer|Maximum number of persists that can be pending but not started. If this limit would be exceeded by a new intermediate persist, ingestion will block until the currently-running persist finishes. Maximum heap memory usage for indexing scales with maxRowsInMemory * (2 + maxPendingPersists).|no (default == 0, meaning one persist can be running concurrently with ingestion, and none can be queued up)| +|`indexSpec`|Object|Tune how data is indexed, see 'IndexSpec' below for more details.|no| +|`reportParseExceptions`|Boolean|If true, exceptions encountered during parsing will be thrown and will halt ingestion; if false, unparseable rows and fields will be skipped.|no (default == false)| +|`handoffConditionTimeout`|Long|Milliseconds to wait for segment handoff. It must be >= 0, where 0 means to wait forever.|no (default == 0)| +|`resetOffsetAutomatically`|Boolean|Whether to reset the consumer sequence numbers if the next sequence number that it is trying to fetch is less than the earliest available sequence number for that particular shard. The sequence number will be reset to either the earliest or latest sequence number depending on `useEarliestOffset` property of `KinesisSupervisorIOConfig` (see below). This situation typically occurs when messages in Kinesis are no longer available for consumption and therefore won't be ingested into Druid. If set to false then ingestion for that particular shard will halt and manual intervention is required to correct the situation, please see `Reset Supervisor` API below.|no (default == false)| +|`skipSequenceNumberAvailabilityCheck`|Boolean|Whether to enable checking if the current sequence number is still available in a particular Kinesis shard. If set to false, the indexing task will attempt to reset the current sequence number (or not), depending on the value of `resetOffsetAutomatically`. |no (default == false)| +|`workerThreads`|Integer|The number of threads that will be used by the supervisor for asynchronous operations.|no (default == min(10, taskCount))| +|`chatThreads`|Integer|The number of threads that will be used for communicating with indexing tasks.|no (default == min(10, taskCount * replicas))| +|`chatRetries`|Integer|The number of times HTTP requests to indexing tasks will be retried before considering tasks unresponsive.|no (default == 8)| +|`httpTimeout`|ISO8601 Period|How long to wait for a HTTP response from an indexing task.|no (default == PT10S)| +|`shutdownTimeout`|ISO8601 Period|How long to wait for the supervisor to attempt a graceful shutdown of tasks before exiting.|no (default == PT80S)| +|`recordBufferSize`|Integer|Size of the buffer (number of events) used between the Kinesis fetch threads and the main ingestion thread.|no (default == 10000)| +|`recordBufferOfferTimeout`|Integer|Length of time in milliseconds to wait for space to become available in the buffer before timing out.|no (default == 5000)| +|`recordBufferFullWait`|Integer|Length of time in milliseconds to wait for the buffer to drain before attempting to fetch records from Kinesis again.|no (default == 5000)| +|`fetchSequenceNumberTimeout`|Integer|Length of time in milliseconds to wait for Kinesis to return the earliest or latest sequence number for a shard. Kinesis will not return the latest sequence number if no data is actively being written to that shard. In this case, this fetch call will repeatedly timeout and retry until fresh data is written to the stream.|no (default == 60000)| +|`fetchThreads`|Integer|Size of the pool of threads fetching data from Kinesis. There is no benefit in having more threads than Kinesis shards.|no (default == max(1, {numProcessors} - 1))| +|`segmentWriteOutMediumFactory`|Object|Segment write-out medium to use when creating segments. See below for more information.|no (not specified by default, the value from `druid.peon.defaultSegmentWriteOutMediumFactory.type` is used)| +|`intermediateHandoffPeriod`|ISO8601 Period|How often the tasks should hand off segments. Handoff will happen either if `maxRowsPerSegment` or `maxTotalRows` is hit or every `intermediateHandoffPeriod`, whichever happens earlier.|no (default == P2147483647D)| +|`logParseExceptions`|Boolean|If true, log an error message when a parsing exception occurs, containing information about the row where the error occurred.|no, default == false| +|`maxParseExceptions`|Integer|The maximum number of parse exceptions that can occur before the task halts ingestion and fails. Overridden if `reportParseExceptions` is set.|no, unlimited default| +|`maxSavedParseExceptions`|Integer|When a parse exception occurs, Druid can keep track of the most recent parse exceptions. "maxSavedParseExceptions" limits how many exception instances will be saved. These saved exceptions will be made available after the task finishes in the [task completion report](../../ingestion/reports.html). Overridden if `reportParseExceptions` is set.|no, default == 0| +|`maxRecordsPerPoll`|Integer| The maximum number of records/events to be fetched from buffer per poll. The actual maximum will be `Max(maxRecordsPerPoll, Max(bufferSize, 1)) |no, default == 100| + +#### IndexSpec + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|bitmap|Object|Compression format for bitmap indexes. Should be a JSON object; see below for options.|no (defaults to Concise)| +|dimensionCompression|String|Compression format for dimension columns. Choose from `LZ4`, `LZF`, or `uncompressed`.|no (default == `LZ4`)| +|metricCompression|String|Compression format for metric columns. Choose from `LZ4`, `LZF`, `uncompressed`, or `none`.|no (default == `LZ4`)| +|longEncoding|String|Encoding format for metric and dimension columns with type long. Choose from `auto` or `longs`. `auto` encodes the values using sequence number or lookup table depending on column cardinality, and store them with variable size. `longs` stores the value as is with 8 bytes each.|no (default == `longs`)| + +##### Bitmap types + +For Concise bitmaps: + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|`type`|String|Must be `concise`.|yes| + +For Roaring bitmaps: + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|`type`|String|Must be `roaring`.|yes| +|`compressRunOnSerialization`|Boolean|Use a run-length encoding where it is estimated as more space efficient.|no (default == `true`)| + +#### SegmentWriteOutMediumFactory + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|`type`|String|See [Additional Peon Configuration: SegmentWriteOutMediumFactory](../../configuration/index.html#segmentwriteoutmediumfactory) for explanation and available options.|yes| + +### KinesisSupervisorIOConfig + +|Field|Type|Description|Required| +|-----|----|-----------|--------| +|`stream`|String|The Kinesis stream to read.|yes| +|`endpoint`|String|The AWS Kinesis stream endpoint for a region. You can find a list of endpoints [here](http://docs.aws.amazon.com/general/latest/gr/rande.html#ak_region).|no (default == kinesis.us-east-1.amazonaws.com)| +|`replicas`|Integer|The number of replica sets, where 1 means a single set of tasks (no replication). Replica tasks will always be assigned to different workers to provide resiliency against node failure.|no (default == 1)| +|`taskCount`|Integer|The maximum number of *reading* tasks in a *replica set*. This means that the maximum number of reading tasks will be `taskCount * replicas` and the total number of tasks (*reading* + *publishing*) will be higher than this. See 'Capacity Planning' below for more details. The number of reading tasks will be less than `taskCount` if `taskCount > {numKinesisshards}`.|no (default == 1)| +|`taskDuration`|ISO8601 Period|The length of time before tasks stop reading and begin publishing their segment.|no (default == PT1H)| +|`startDelay`|ISO8601 Period|The period to wait before the supervisor starts managing tasks.|no (default == PT5S)| +|`period`|ISO8601 Period|How often the supervisor will execute its management logic. Note that the supervisor will also run in response to certain events (such as tasks succeeding, failing, and reaching their taskDuration) so this value specifies the maximum time between iterations.|no (default == PT30S)| +|`useEarliestSequenceNumber`|Boolean|If a supervisor is managing a dataSource for the first time, it will obtain a set of starting sequence numbers from Kinesis. This flag determines whether it retrieves the earliest or latest sequence numbers in Kinesis. Under normal circumstances, subsequent tasks will start from where the previous segments ended so this flag will only be used on first run.|no (default == false)| +|`completionTimeout`|ISO8601 Period|The length of time to wait before declaring a publishing task as failed and terminating it. If this is set too low, your tasks may never publish. The publishing clock for a task begins roughly after `taskDuration` elapses.|no (default == PT6H)| +|`lateMessageRejectionPeriod`|ISO8601 Period|Configure tasks to reject messages with timestamps earlier than this period before the task was created; for example if this is set to `PT1H` and the supervisor creates a task at *2016-01-01T12:00Z*, messages with timestamps earlier than *2016-01-01T11:00Z* will be dropped. This may help prevent concurrency issues if your data stream has late messages and you have multiple pipelines that need to operate on the same segments (e.g. a realtime and a nightly batch ingestion pipeline).|no (default == none)| +|`earlyMessageRejectionPeriod`|ISO8601 Period|Configure tasks to reject messages with timestamps later than this period after the task reached its taskDuration; for example if this is set to `PT1H`, the taskDuration is set to `PT1H` and the supervisor creates a task at *2016-01-01T12:00Z*, messages with timestamps later than *2016-01-01T14:00Z* will be dropped. **Note:** Tasks sometimes run past their task duration, for example, in cases of supervisor failover. Setting earlyMessageRejectionPeriod too low may cause messages to be dropped unexpectedly whenever a task runs past its originally configured task duration.|no (default == none)| +|`recordsPerFetch`|Integer|The number of records to request per GetRecords call to Kinesis. See 'Determining Fetch Settings' below.|no (default == 2000)| +|`fetchDelayMillis`|Integer|Time in milliseconds to wait between subsequent GetRecords calls to Kinesis. See 'Determining Fetch Settings' below.|no (default == 1000)| +|`awsAssumedRoleArn`|String|The AWS assumed role to use for additional permissions.|no| +|`awsExternalId`|String|The AWS external id to use for additional permissions.|no| +|`deaggregate`|Boolean|Whether to use the de-aggregate function of the KCL. See below for details.|no| + +## Operations + +This section gives descriptions of how some supervisor APIs work specifically in Kinesis Indexing Service. +For all supervisor APIs, please check [Supervisor APIs](../../operations/api-reference.html#supervisors). + +### AWS Authentication +To authenticate with AWS, you must provide your AWS access key and AWS secret key via runtime.properties, for example: +``` +-Ddruid.kinesis.accessKey=123 -Ddruid.kinesis.secretKey=456 +``` +The AWS access key ID and secret access key are used for Kinesis API requests. If this is not provided, the service will look for credentials set in environment variables, in the default profile configuration file, and from the EC2 instance profile provider (in this order). + +### Getting Supervisor Status Report + +`GET /druid/indexer/v1/supervisor//status` returns a snapshot report of the current state of the tasks managed by the given supervisor. This includes the latest +sequence numbers as reported by Kinesis. Unlike the Kafka Indexing Service, stats about lag is not yet supported. + +### Updating Existing Supervisors + +`POST /druid/indexer/v1/supervisor` can be used to update existing supervisor spec. +Calling this endpoint when there is already an existing supervisor for the same dataSource will cause: + +- The running supervisor to signal its managed tasks to stop reading and begin publishing. +- The running supervisor to exit. +- A new supervisor to be created using the configuration provided in the request body. This supervisor will retain the +existing publishing tasks and will create new tasks starting at the sequence numbers the publishing tasks ended on. + +Seamless schema migrations can thus be achieved by simply submitting the new schema using this endpoint. + +### Suspending and Resuming Supervisors + +You can suspend and resume a supervisor using `POST /druid/indexer/v1/supervisor//suspend` and `POST /druid/indexer/v1/supervisor//resume`, respectively. + +Note that the supervisor itself will still be operating and emitting logs and metrics, +it will just ensure that no indexing tasks are running until the supervisor is resumed. + +### Resetting Supervisors + +To reset a running supervisor, you can use `POST /druid/indexer/v1/supervisor//reset`. + +The indexing service keeps track of the latest persisted Kinesis sequence number in order to provide exactly-once ingestion +guarantees across tasks. Subsequent tasks must start reading from where the previous task completed in order for the +generated segments to be accepted. If the messages at the expected starting sequence numbers are no longer available in Kinesis +(typically because the message retention period has elapsed or the topic was removed and re-created) the supervisor will +refuse to start and in-flight tasks will fail. + +This endpoint can be used to clear the stored sequence numbers which will cause the supervisor to start reading from +either the earliest or latest sequence numbers in Kinesis (depending on the value of `useEarliestSequenceNumber`). The supervisor must be +running for this endpoint to be available. After the stored sequence numbers are cleared, the supervisor will automatically kill +and re-create any active tasks so that tasks begin reading from valid sequence numbers. + +Note that since the stored sequence numbers are necessary to guarantee exactly-once ingestion, resetting them with this endpoint +may cause some Kinesis messages to be skipped or to be read twice. + +### Terminating Supervisors + +`POST /druid/indexer/v1/supervisor//terminate` terminates a supervisor and causes all associated indexing +tasks managed by this supervisor to immediately stop and begin +publishing their segments. This supervisor will still exist in the metadata store and it's history may be retrieved +with the supervisor history api, but will not be listed in the 'get supervisors' api response nor can it's configuration +or status report be retrieved. The only way this supervisor can start again is by submitting a functioning supervisor +spec to the create api. + +### Capacity Planning + +Kinesis indexing tasks run on middle managers and are thus limited by the resources available in the middle manager +cluster. In particular, you should make sure that you have sufficient worker capacity (configured using the +`druid.worker.capacity` property) to handle the configuration in the supervisor spec. Note that worker capacity is +shared across all types of indexing tasks, so you should plan your worker capacity to handle your total indexing load +(e.g. batch processing, realtime tasks, merging tasks, etc.). If your workers run out of capacity, Kinesis indexing tasks +will queue and wait for the next available worker. This may cause queries to return partial results but will not result +in data loss (assuming the tasks run before Kinesis purges those sequence numbers). + +A running task will normally be in one of two states: *reading* or *publishing*. A task will remain in reading state for +`taskDuration`, at which point it will transition to publishing state. A task will remain in publishing state for as long +as it takes to generate segments, push segments to deep storage, and have them be loaded and served by a historical node +(or until `completionTimeout` elapses). + +The number of reading tasks is controlled by `replicas` and `taskCount`. In general, there will be `replicas * taskCount` +reading tasks, the exception being if taskCount > {numKinesisShards} in which case {numKinesisShards} tasks will +be used instead. When `taskDuration` elapses, these tasks will transition to publishing state and `replicas * taskCount` +new reading tasks will be created. Therefore to allow for reading tasks and publishing tasks to run concurrently, there +should be a minimum capacity of: + +``` +workerCapacity = 2 * replicas * taskCount +``` + +This value is for the ideal situation in which there is at most one set of tasks publishing while another set is reading. +In some circumstances, it is possible to have multiple sets of tasks publishing simultaneously. This would happen if the +time-to-publish (generate segment, push to deep storage, loaded on historical) > `taskDuration`. This is a valid +scenario (correctness-wise) but requires additional worker capacity to support. In general, it is a good idea to have +`taskDuration` be large enough that the previous set of tasks finishes publishing before the current set begins. + +### Supervisor Persistence + +When a supervisor spec is submitted via the `POST /druid/indexer/v1/supervisor` endpoint, it is persisted in the +configured metadata database. There can only be a single supervisor per dataSource, and submitting a second spec for +the same dataSource will overwrite the previous one. + +When an overlord gains leadership, either by being started or as a result of another overlord failing, it will spawn +a supervisor for each supervisor spec in the metadata database. The supervisor will then discover running Kinesis indexing +tasks and will attempt to adopt them if they are compatible with the supervisor's configuration. If they are not +compatible because they have a different ingestion spec or shard allocation, the tasks will be killed and the +supervisor will create a new set of tasks. In this way, the supervisors are persistent across overlord restarts and +fail-overs. + +A supervisor is stopped via the `POST /druid/indexer/v1/supervisor//terminate` endpoint. This places a +tombstone marker in the database (to prevent the supervisor from being reloaded on a restart) and then gracefully +shuts down the currently running supervisor. When a supervisor is shut down in this way, it will instruct its +managed tasks to stop reading and begin publishing their segments immediately. The call to the shutdown endpoint will +return after all tasks have been signalled to stop but before the tasks finish publishing their segments. + +### Schema/Configuration Changes + +Schema and configuration changes are handled by submitting the new supervisor spec via the same +`POST /druid/indexer/v1/supervisor` endpoint used to initially create the supervisor. The overlord will initiate a +graceful shutdown of the existing supervisor which will cause the tasks being managed by that supervisor to stop reading +and begin publishing their segments. A new supervisor will then be started which will create a new set of tasks that +will start reading from the sequence numbers where the previous now-publishing tasks left off, but using the updated schema. +In this way, configuration changes can be applied without requiring any pause in ingestion. + +### Deployment Notes + +#### On the Subject of Segments + +Each Kinesis Indexing Task puts events consumed from Kinesis Shards assigned to it in a single segment for each segment +granular interval until maxRowsPerSegment, maxTotalRows or intermediateHandoffPeriod limit is reached, at this point a new shard +for this segment granularity is created for further events. Kinesis Indexing Task also does incremental hand-offs which +means that all the segments created by a task will not be held up till the task duration is over. As soon as maxRowsPerSegment, +maxTotalRows or intermediateHandoffPeriod limit is hit, all the segments held by the task at that point in time will be handed-off +and new set of segments will be created for further events. This means that the task can run for longer durations of time +without accumulating old segments locally on Middle Manager nodes and it is encouraged to do so. + +Kinesis Indexing Service may still produce some small segments. Lets say the task duration is 4 hours, segment granularity +is set to an HOUR and Supervisor was started at 9:10 then after 4 hours at 13:10, new set of tasks will be started and +events for the interval 13:00 - 14:00 may be split across previous and new set of tasks. If you see it becoming a problem then +one can schedule re-indexing tasks be run to merge segments together into new segments of an ideal size (in the range of ~500-700 MB per segment). +Details on how to optimize the segment size can be found on [Segment size optimization](../../operations/segment-optimization.html). +There is also ongoing work to support automatic segment compaction of sharded segments as well as compaction not requiring +Hadoop (see [here](https://github.com/apache/incubator-druid/pull/5102)). + +### Determining Fetch Settings +Internally, the Kinesis Indexing Service uses the Kinesis Record Supplier abstraction for fetching Kinesis data records and storing the records +locally. The way the Kinesis Record Supplier fetches records is to have a separate thread run the fetching operation per each Kinesis Shard, the +max number of threads is determined by `fetchThreads`. For example, a Kinesis stream with 3 shards will have 3 threads, each fetching from a shard separately. +There is a delay between each fetching operation, which is controlled by `fetchDelayMillis`. The maximum number of records to be fetched per thread per +operation is controlled by `recordsPerFetch`. Note that this is not the same as `maxRecordsPerPoll`. + +The records fetched by each thread will be pushed to a queue in the order that they are fetched. The records are stored in this queue until `poll()` is called +by either the supervisor or the indexing task. `poll()` will attempt to drain the internal buffer queue up to a limit of `max(maxRecordsPerPoll, q.size())`. +Here `maxRecordsPerPoll` controls the theoretical maximum records to drain out of the buffer queue, so setting this parameter to a reasonable value is essential +in preventing the queue from overflowing or memory exceeding heap size. + +Kinesis places the following restrictions on calls to fetch records: + +- Each data record can be up to 1 MB in size. +- Each shard can support up to 5 transactions per second for reads. +- Each shard can read up to 2 MB per second. +- The maximum size of data that GetRecords can return is 10 MB. + +Values for `recordsPerFetch` and `fetchDelayMillis` should be chosen to maximize throughput under the above constraints. +The values that you choose will depend on the average size of a record and the number of consumers you have reading from +a given shard (which will be `replicas` unless you have other consumers also reading from this Kinesis stream). + +If the above limits are violated, AWS will throw ProvisionedThroughputExceededException errors on subsequent calls to +read data. When this happens, the Kinesis indexing service will pause by `fetchDelayMillis` and then attempt the call +again. + +Internally, each indexing task maintains a buffer that stores the fetched but not yet processed record. `recordsPerFetch` and `fetchDelayMillis` +control this behavior. The number of records that the indexing task fetch from the buffer is controlled by `maxRecordsPerPoll`, which +determines the number of records to be processed per each ingestion loop in the task. + +## Deaggregation +See [issue](https://github.com/apache/incubator-druid/issues/6714) + +The Kinesis indexing service supports de-aggregation of multiple rows packed into a single record by the Kinesis +Producer Library's aggregate method for more efficient data transfer. Currently, enabling the de-aggregate functionality +requires the user to manually provide the Kinesis Client Library on the classpath, since this library has a license not +compatible with Apache projects. + +To enable this feature, add the `amazon-kinesis-client` (tested on version `1.9.2`) jar file ([link](https://mvnrepository.com/artifact/com.amazonaws/amazon-kinesis-client/1.9.2)) under `dist/druid/extensions/druid-kinesis-indexing-service/`. +Then when submitting a supervisor-spec, set `deaggregate` to true. \ No newline at end of file diff --git a/docs/content/development/extensions.md b/docs/content/development/extensions.md index 1f9d925dad9c..a5b8a5b3f57a 100644 --- a/docs/content/development/extensions.md +++ b/docs/content/development/extensions.md @@ -51,6 +51,7 @@ Core extensions are maintained by Druid committers. |druid-kafka-eight|Kafka ingest firehose (high level consumer) for realtime nodes.|[link](../development/extensions-core/kafka-eight-firehose.html)| |druid-kafka-extraction-namespace|Kafka-based namespaced lookup. Requires namespace lookup extension.|[link](../development/extensions-core/kafka-extraction-namespace.html)| |druid-kafka-indexing-service|Supervised exactly-once Kafka ingestion for the indexing service.|[link](../development/extensions-core/kafka-ingestion.html)| +|druid-kinesis-indexing-service|Supervised exactly-once Kinesis ingestion for the indexing service.|[link](../development/extensions-core/kinesis-ingestion.html)| |druid-kerberos|Kerberos authentication for druid nodes.|[link](../development/extensions-core/druid-kerberos.html)| |druid-lookups-cached-global|A module for [lookups](../querying/lookups.html) providing a jvm-global eager caching for lookups. It provides JDBC and URI implementations for fetching lookup data.|[link](../development/extensions-core/lookups-cached-global.html)| |druid-lookups-cached-single| Per lookup caching module to support the use cases where a lookup need to be isolated from the global pool of lookups |[link](../development/extensions-core/druid-lookups.html)| diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/IncrementalPublishingKafkaIndexTaskRunner.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/IncrementalPublishingKafkaIndexTaskRunner.java index c05021df7bbd..e8f626208667 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/IncrementalPublishingKafkaIndexTaskRunner.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/IncrementalPublishingKafkaIndexTaskRunner.java @@ -19,185 +19,48 @@ package org.apache.druid.indexing.kafka; -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; -import com.google.common.base.Function; -import com.google.common.base.Joiner; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Optional; -import com.google.common.base.Preconditions; -import com.google.common.base.Supplier; -import com.google.common.base.Throwables; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Sets; -import com.google.common.primitives.Longs; -import com.google.common.util.concurrent.FutureCallback; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.SettableFuture; -import org.apache.druid.data.input.Committer; -import org.apache.druid.data.input.InputRow; import org.apache.druid.data.input.impl.InputRowParser; -import org.apache.druid.discovery.DiscoveryDruidNode; -import org.apache.druid.discovery.LookupNodeService; -import org.apache.druid.discovery.NodeType; -import org.apache.druid.indexer.IngestionState; -import org.apache.druid.indexer.TaskStatus; -import org.apache.druid.indexing.common.IngestionStatsAndErrorsTaskReport; -import org.apache.druid.indexing.common.IngestionStatsAndErrorsTaskReportData; -import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder; -import org.apache.druid.indexing.common.TaskReport; import org.apache.druid.indexing.common.TaskToolbox; -import org.apache.druid.indexing.common.actions.CheckPointDataSourceMetadataAction; -import org.apache.druid.indexing.common.actions.ResetDataSourceMetadataAction; -import org.apache.druid.indexing.common.actions.SegmentTransactionalInsertAction; -import org.apache.druid.indexing.common.stats.RowIngestionMeters; import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; -import org.apache.druid.indexing.common.task.IndexTaskUtils; -import org.apache.druid.indexing.common.task.RealtimeIndexTask; -import org.apache.druid.indexing.kafka.KafkaIndexTask.Status; -import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; import org.apache.druid.java.util.common.ISE; -import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.collect.Utils; -import org.apache.druid.java.util.common.parsers.ParseException; import org.apache.druid.java.util.emitter.EmittingLogger; -import org.apache.druid.segment.indexing.RealtimeIOConfig; -import org.apache.druid.segment.realtime.FireDepartment; -import org.apache.druid.segment.realtime.FireDepartmentMetrics; -import org.apache.druid.segment.realtime.appenderator.Appenderator; -import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; -import org.apache.druid.segment.realtime.appenderator.SegmentsAndMetadata; -import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; -import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; -import org.apache.druid.server.security.Access; -import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.timeline.DataSegment; import org.apache.druid.utils.CircularBuffer; -import org.apache.kafka.clients.consumer.ConsumerRecord; -import org.apache.kafka.clients.consumer.ConsumerRecords; -import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.OffsetOutOfRangeException; import org.apache.kafka.common.TopicPartition; -import org.joda.time.DateTime; +import javax.annotation.Nonnull; import javax.annotation.Nullable; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.Consumes; -import javax.ws.rs.DefaultValue; -import javax.ws.rs.GET; -import javax.ws.rs.POST; -import javax.ws.rs.Path; -import javax.ws.rs.Produces; -import javax.ws.rs.QueryParam; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.MediaType; -import javax.ws.rs.core.Response; -import java.io.File; +import javax.validation.constraints.NotNull; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; -import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.Map.Entry; import java.util.Set; import java.util.TreeMap; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Future; -import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.locks.Condition; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; /** * Kafka indexing task runner supporting incremental segments publishing */ -public class IncrementalPublishingKafkaIndexTaskRunner implements KafkaIndexTaskRunner +public class IncrementalPublishingKafkaIndexTaskRunner extends SeekableStreamIndexTaskRunner { private static final EmittingLogger log = new EmittingLogger(IncrementalPublishingKafkaIndexTaskRunner.class); - private static final String METADATA_NEXT_PARTITIONS = "nextPartitions"; - private static final String METADATA_PUBLISH_PARTITIONS = "publishPartitions"; - - private final Map endOffsets; - private final Map nextOffsets = new ConcurrentHashMap<>(); - private final Map lastPersistedOffsets = new ConcurrentHashMap<>(); - - // The pause lock and associated conditions are to support coordination between the Jetty threads and the main - // ingestion loop. The goal is to provide callers of the API a guarantee that if pause() returns successfully - // the ingestion loop has been stopped at the returned offsets and will not ingest any more data until resumed. The - // fields are used as follows (every step requires acquiring [pauseLock]): - // Pausing: - // - In pause(), [pauseRequested] is set to true and then execution waits for [status] to change to PAUSED, with the - // condition checked when [hasPaused] is signalled. - // - In possiblyPause() called from the main loop, if [pauseRequested] is true, [status] is set to PAUSED, - // [hasPaused] is signalled, and execution pauses until [pauseRequested] becomes false, either by being set or by - // the [pauseMillis] timeout elapsing. [pauseRequested] is checked when [shouldResume] is signalled. - // Resuming: - // - In resume(), [pauseRequested] is set to false, [shouldResume] is signalled, and execution waits for [status] to - // change to something other than PAUSED, with the condition checked when [shouldResume] is signalled. - // - In possiblyPause(), when [shouldResume] is signalled, if [pauseRequested] has become false the pause loop ends, - // [status] is changed to STARTING and [shouldResume] is signalled. - - private final Lock pauseLock = new ReentrantLock(); - private final Condition hasPaused = pauseLock.newCondition(); - private final Condition shouldResume = pauseLock.newCondition(); - - private final AtomicBoolean stopRequested = new AtomicBoolean(false); - private final AtomicBoolean publishOnStop = new AtomicBoolean(false); - - // [statusLock] is used to synchronize the Jetty thread calling stopGracefully() with the main run thread. It prevents - // the main run thread from switching into a publishing state while the stopGracefully() thread thinks it's still in - // a pre-publishing state. This is important because stopGracefully() will try to use the [stopRequested] flag to stop - // the main thread where possible, but this flag is not honored once publishing has begun so in this case we must - // interrupt the thread. The lock ensures that if the run thread is about to transition into publishing state, it - // blocks until after stopGracefully() has set [stopRequested] and then does a final check on [stopRequested] before - // transitioning to publishing state. - private final Object statusLock = new Object(); - - private final Lock pollRetryLock = new ReentrantLock(); - private final Condition isAwaitingRetry = pollRetryLock.newCondition(); - private final KafkaIndexTask task; - private final KafkaIOConfig ioConfig; - private final KafkaTuningConfig tuningConfig; - private final InputRowParser parser; - private final AuthorizerMapper authorizerMapper; - private final Optional chatHandlerProvider; - private final CircularBuffer savedParseExceptions; - private final String topic; - private final RowIngestionMeters rowIngestionMeters; - - private final Set publishingSequences = Sets.newConcurrentHashSet(); - private final List> publishWaitList = new ArrayList<>(); - private final List> handOffWaitList = new ArrayList<>(); - - private volatile DateTime startTime; - private volatile Status status = Status.NOT_STARTED; // this is only ever set by the task runner thread (runThread) - private volatile TaskToolbox toolbox; - private volatile Thread runThread; - private volatile Appenderator appenderator; - private volatile StreamAppenderatorDriver driver; - private volatile FireDepartmentMetrics fireDepartmentMetrics; - private volatile IngestionState ingestionState; - - private volatile boolean pauseRequested = false; - private volatile long nextCheckpointTime; - - private volatile CopyOnWriteArrayList sequences; - private volatile Throwable backgroundThreadException; public IncrementalPublishingKafkaIndexTaskRunner( KafkaIndexTask task, @@ -208,854 +71,85 @@ public IncrementalPublishingKafkaIndexTaskRunner( RowIngestionMetersFactory rowIngestionMetersFactory ) { + super( + task, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory + ); this.task = task; - this.ioConfig = task.getIOConfig(); - this.tuningConfig = task.getTuningConfig(); - this.parser = parser; - this.authorizerMapper = authorizerMapper; - this.chatHandlerProvider = chatHandlerProvider; - this.savedParseExceptions = savedParseExceptions; - this.topic = ioConfig.getStartPartitions().getTopic(); - this.rowIngestionMeters = rowIngestionMetersFactory.createRowIngestionMeters(); - - this.endOffsets = new ConcurrentHashMap<>(ioConfig.getEndPartitions().getPartitionOffsetMap()); - this.sequences = new CopyOnWriteArrayList<>(); - this.ingestionState = IngestionState.NOT_STARTED; - - resetNextCheckpointTime(); } @Override - public TaskStatus run(TaskToolbox toolbox) + protected Long getSequenceNumberToStoreAfterRead(@NotNull Long sequenceNumber) { - try { - return runInternal(toolbox); - } - catch (Exception e) { - log.error(e, "Encountered exception while running task."); - final String errorMsg = Throwables.getStackTraceAsString(e); - toolbox.getTaskReportFileWriter().write(getTaskCompletionReports(errorMsg)); - return TaskStatus.failure( - task.getId(), - errorMsg - ); - } + return sequenceNumber + 1; } - private TaskStatus runInternal(TaskToolbox toolbox) throws Exception - { - log.info("Starting up!"); - - startTime = DateTimes.nowUtc(); - status = Status.STARTING; - this.toolbox = toolbox; - - if (!restoreSequences()) { - final TreeMap> checkpoints = getCheckPointsFromContext(toolbox, task); - if (checkpoints != null) { - Iterator>> sequenceOffsets = checkpoints.entrySet().iterator(); - Map.Entry> previous = sequenceOffsets.next(); - while (sequenceOffsets.hasNext()) { - Map.Entry> current = sequenceOffsets.next(); - sequences.add(new SequenceMetadata( - previous.getKey(), - StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), previous.getKey()), - previous.getValue(), - current.getValue(), - true - )); - previous = current; - } - sequences.add(new SequenceMetadata( - previous.getKey(), - StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), previous.getKey()), - previous.getValue(), - endOffsets, - false - )); - } else { - sequences.add(new SequenceMetadata( - 0, - StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), 0), - ioConfig.getStartPartitions().getPartitionOffsetMap(), - endOffsets, - false - )); - } - } - log.info("Starting with sequences: %s", sequences); - - if (chatHandlerProvider.isPresent()) { - log.info("Found chat handler of class[%s]", chatHandlerProvider.get().getClass().getName()); - chatHandlerProvider.get().register(task.getId(), this, false); - } else { - log.warn("No chat handler detected"); - } - - runThread = Thread.currentThread(); - - // Set up FireDepartmentMetrics - final FireDepartment fireDepartmentForMetrics = new FireDepartment( - task.getDataSchema(), - new RealtimeIOConfig(null, null, null), - null - ); - fireDepartmentMetrics = fireDepartmentForMetrics.getMetrics(); - toolbox.getMonitorScheduler() - .addMonitor(TaskRealtimeMetricsMonitorBuilder.build(task, fireDepartmentForMetrics, rowIngestionMeters)); - - final String lookupTier = task.getContextValue(RealtimeIndexTask.CTX_KEY_LOOKUP_TIER); - LookupNodeService lookupNodeService = lookupTier == null ? - toolbox.getLookupNodeService() : - new LookupNodeService(lookupTier); - DiscoveryDruidNode discoveryDruidNode = new DiscoveryDruidNode( - toolbox.getDruidNode(), - NodeType.PEON, - ImmutableMap.of( - toolbox.getDataNodeService().getName(), toolbox.getDataNodeService(), - lookupNodeService.getName(), lookupNodeService - ) - ); - - Throwable caughtExceptionOuter = null; - try (final KafkaConsumer consumer = task.newConsumer()) { - toolbox.getDataSegmentServerAnnouncer().announce(); - toolbox.getDruidNodeAnnouncer().announce(discoveryDruidNode); - - appenderator = task.newAppenderator(fireDepartmentMetrics, toolbox); - driver = task.newDriver(appenderator, toolbox, fireDepartmentMetrics); - - final String topic = ioConfig.getStartPartitions().getTopic(); - - // Start up, set up initial offsets. - final Object restoredMetadata = driver.startJob(); - if (restoredMetadata == null) { - // no persist has happened so far - // so either this is a brand new task or replacement of a failed task - Preconditions.checkState(sequences.get(0).startOffsets.entrySet().stream().allMatch( - partitionOffsetEntry -> Longs.compare( - partitionOffsetEntry.getValue(), - ioConfig.getStartPartitions() - .getPartitionOffsetMap() - .get(partitionOffsetEntry.getKey()) - ) >= 0 - ), "Sequence offsets are not compatible with start offsets of task"); - nextOffsets.putAll(sequences.get(0).startOffsets); - } else { - final Map restoredMetadataMap = (Map) restoredMetadata; - final KafkaPartitions restoredNextPartitions = toolbox.getObjectMapper().convertValue( - restoredMetadataMap.get(METADATA_NEXT_PARTITIONS), - KafkaPartitions.class - ); - nextOffsets.putAll(restoredNextPartitions.getPartitionOffsetMap()); - - // Sanity checks. - if (!restoredNextPartitions.getTopic().equals(ioConfig.getStartPartitions().getTopic())) { - throw new ISE( - "WTF?! Restored topic[%s] but expected topic[%s]", - restoredNextPartitions.getTopic(), - ioConfig.getStartPartitions().getTopic() - ); - } - - if (!nextOffsets.keySet().equals(ioConfig.getStartPartitions().getPartitionOffsetMap().keySet())) { - throw new ISE( - "WTF?! Restored partitions[%s] but expected partitions[%s]", - nextOffsets.keySet(), - ioConfig.getStartPartitions().getPartitionOffsetMap().keySet() - ); - } - // sequences size can be 0 only when all sequences got published and task stopped before it could finish - // which is super rare - if (sequences.size() == 0 || sequences.get(sequences.size() - 1).isCheckpointed()) { - this.endOffsets.putAll(sequences.size() == 0 - ? nextOffsets - : sequences.get(sequences.size() - 1).getEndOffsets()); - log.info("End offsets changed to [%s]", endOffsets); - } - } - - // Set up committer. - final Supplier committerSupplier = () -> { - final Map snapshot = ImmutableMap.copyOf(nextOffsets); - lastPersistedOffsets.clear(); - lastPersistedOffsets.putAll(snapshot); - - return new Committer() - { - @Override - public Object getMetadata() - { - return ImmutableMap.of( - METADATA_NEXT_PARTITIONS, new KafkaPartitions( - ioConfig.getStartPartitions().getTopic(), - snapshot - ) - ); - } - - @Override - public void run() - { - // Do nothing. - } - }; - }; - - // restart publishing of sequences (if any) - maybePersistAndPublishSequences(committerSupplier); - - Set assignment = assignPartitionsAndSeekToNext(consumer, topic); - - ingestionState = IngestionState.BUILD_SEGMENTS; - - // Main loop. - // Could eventually support leader/follower mode (for keeping replicas more in sync) - boolean stillReading = !assignment.isEmpty(); - status = Status.READING; - Throwable caughtExceptionInner = null; - try { - while (stillReading) { - if (possiblyPause()) { - // The partition assignments may have changed while paused by a call to setEndOffsets() so reassign - // partitions upon resuming. This is safe even if the end offsets have not been modified. - assignment = assignPartitionsAndSeekToNext(consumer, topic); - - if (assignment.isEmpty()) { - log.info("All partitions have been fully read"); - publishOnStop.set(true); - stopRequested.set(true); - } - } - - // if stop is requested or task's end offset is set by call to setEndOffsets method with finish set to true - if (stopRequested.get() || sequences.get(sequences.size() - 1).isCheckpointed()) { - status = Status.PUBLISHING; - } - - if (stopRequested.get()) { - break; - } - - if (backgroundThreadException != null) { - throw new RuntimeException(backgroundThreadException); - } - - checkPublishAndHandoffFailure(); - - maybePersistAndPublishSequences(committerSupplier); - - // The retrying business is because the KafkaConsumer throws OffsetOutOfRangeException if the seeked-to - // offset is not present in the topic-partition. This can happen if we're asking a task to read from data - // that has not been written yet (which is totally legitimate). So let's wait for it to show up. - ConsumerRecords records = ConsumerRecords.empty(); - try { - records = consumer.poll(KafkaIndexTask.POLL_TIMEOUT_MILLIS); - } - catch (OffsetOutOfRangeException e) { - log.warn("OffsetOutOfRangeException with message [%s]", e.getMessage()); - possiblyResetOffsetsOrWait(e.offsetOutOfRangePartitions(), consumer, toolbox); - stillReading = !assignment.isEmpty(); - } - - SequenceMetadata sequenceToCheckpoint = null; - for (ConsumerRecord record : records) { - log.trace( - "Got topic[%s] partition[%d] offset[%,d].", - record.topic(), - record.partition(), - record.offset() - ); - - if (record.offset() < endOffsets.get(record.partition())) { - if (record.offset() != nextOffsets.get(record.partition())) { - if (ioConfig.isSkipOffsetGaps()) { - log.warn( - "Skipped to offset[%,d] after offset[%,d] in partition[%d].", - record.offset(), - nextOffsets.get(record.partition()), - record.partition() - ); - } else { - throw new ISE( - "WTF?! Got offset[%,d] after offset[%,d] in partition[%d].", - record.offset(), - nextOffsets.get(record.partition()), - record.partition() - ); - } - } - - try { - final byte[] valueBytes = record.value(); - final List rows = valueBytes == null - ? Utils.nullableListOf((InputRow) null) - : parser.parseBatch(ByteBuffer.wrap(valueBytes)); - boolean isPersistRequired = false; - - final SequenceMetadata sequenceToUse = sequences - .stream() - .filter(sequenceMetadata -> sequenceMetadata.canHandle(record)) - .findFirst() - .orElse(null); - - if (sequenceToUse == null) { - throw new ISE( - "WTH?! cannot find any valid sequence for record with partition [%d] and offset [%d]. Current sequences: %s", - record.partition(), - record.offset(), - sequences - ); - } - - for (InputRow row : rows) { - if (row != null && task.withinMinMaxRecordTime(row)) { - final AppenderatorDriverAddResult addResult = driver.add( - row, - sequenceToUse.getSequenceName(), - committerSupplier, - // skip segment lineage check as there will always be one segment - // for combination of sequence and segment granularity. - // It is necessary to skip it as the task puts messages polled from all the - // assigned Kafka partitions into a single Druid segment, thus ordering of - // messages among replica tasks across assigned partitions is not guaranteed - // which may cause replica tasks to ask for segments with different interval - // in different order which might cause SegmentAllocateAction to fail. - true, - // do not allow incremental persists to happen until all the rows from this batch - // of rows are indexed - false - ); - - if (addResult.isOk()) { - // If the number of rows in the segment exceeds the threshold after adding a row, - // move the segment out from the active segments of BaseAppenderatorDriver to make a new segment. - if (addResult.isPushRequired(tuningConfig) && !sequenceToUse.isCheckpointed()) { - sequenceToCheckpoint = sequenceToUse; - } - isPersistRequired |= addResult.isPersistRequired(); - } else { - // Failure to allocate segment puts determinism at risk, bail out to be safe. - // May want configurable behavior here at some point. - // If we allow continuing, then consider blacklisting the interval for a while to avoid constant checks. - throw new ISE("Could not allocate segment for row with timestamp[%s]", row.getTimestamp()); - } - - if (addResult.getParseException() != null) { - handleParseException(addResult.getParseException(), record); - } else { - rowIngestionMeters.incrementProcessed(); - } - } else { - rowIngestionMeters.incrementThrownAway(); - } - } - if (isPersistRequired) { - Futures.addCallback( - driver.persistAsync(committerSupplier.get()), - new FutureCallback() - { - @Override - public void onSuccess(@Nullable Object result) - { - log.info("Persist completed with metadata [%s]", result); - } - - @Override - public void onFailure(Throwable t) - { - log.error("Persist failed, dying"); - backgroundThreadException = t; - } - } - ); - } - } - catch (ParseException e) { - handleParseException(e, record); - } - - nextOffsets.put(record.partition(), record.offset() + 1); - } - - if (nextOffsets.get(record.partition()).equals(endOffsets.get(record.partition())) - && assignment.remove(record.partition())) { - log.info("Finished reading topic[%s], partition[%,d].", record.topic(), record.partition()); - KafkaIndexTask.assignPartitions(consumer, topic, assignment); - stillReading = !assignment.isEmpty(); - } - } - - if (System.currentTimeMillis() > nextCheckpointTime) { - sequenceToCheckpoint = sequences.get(sequences.size() - 1); - } - - if (sequenceToCheckpoint != null && stillReading) { - Preconditions.checkArgument( - sequences.get(sequences.size() - 1) - .getSequenceName() - .equals(sequenceToCheckpoint.getSequenceName()), - "Cannot checkpoint a sequence [%s] which is not the latest one, sequences %s", - sequenceToCheckpoint, - sequences - ); - requestPause(); - final CheckPointDataSourceMetadataAction checkpointAction = new CheckPointDataSourceMetadataAction( - task.getDataSource(), - ioConfig.getTaskGroupId(), - task.getIOConfig().getBaseSequenceName(), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, sequenceToCheckpoint.getStartOffsets())), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, nextOffsets)) - ); - if (!toolbox.getTaskActionClient().submit(checkpointAction)) { - throw new ISE("Checkpoint request with offsets [%s] failed, dying", nextOffsets); - } - } - } - ingestionState = IngestionState.COMPLETED; - } - catch (Exception e) { - // (1) catch all exceptions while reading from kafka - caughtExceptionInner = e; - log.error(e, "Encountered exception in run() before persisting."); - throw e; - } - finally { - log.info("Persisting all pending data"); - try { - driver.persist(committerSupplier.get()); // persist pending data - } - catch (Exception e) { - if (caughtExceptionInner != null) { - caughtExceptionInner.addSuppressed(e); - } else { - throw e; - } - } - } - - synchronized (statusLock) { - if (stopRequested.get() && !publishOnStop.get()) { - throw new InterruptedException("Stopping without publishing"); - } - - status = Status.PUBLISHING; - } - - for (SequenceMetadata sequenceMetadata : sequences) { - if (!publishingSequences.contains(sequenceMetadata.getSequenceName())) { - // this is done to prevent checks in sequence specific commit supplier from failing - sequenceMetadata.setEndOffsets(nextOffsets); - sequenceMetadata.updateAssignments(nextOffsets); - publishingSequences.add(sequenceMetadata.getSequenceName()); - // persist already done in finally, so directly add to publishQueue - publishAndRegisterHandoff(sequenceMetadata); - } - } - - if (backgroundThreadException != null) { - throw new RuntimeException(backgroundThreadException); - } - - // Wait for publish futures to complete. - Futures.allAsList(publishWaitList).get(); - - // Wait for handoff futures to complete. - // Note that every publishing task (created by calling AppenderatorDriver.publish()) has a corresponding - // handoffFuture. handoffFuture can throw an exception if 1) the corresponding publishFuture failed or 2) it - // failed to persist sequences. It might also return null if handoff failed, but was recoverable. - // See publishAndRegisterHandoff() for details. - List handedOffList = Collections.emptyList(); - if (tuningConfig.getHandoffConditionTimeout() == 0) { - handedOffList = Futures.allAsList(handOffWaitList).get(); - } else { - try { - handedOffList = Futures.allAsList(handOffWaitList) - .get(tuningConfig.getHandoffConditionTimeout(), TimeUnit.MILLISECONDS); - } - catch (TimeoutException e) { - // Handoff timeout is not an indexing failure, but coordination failure. We simply ignore timeout exception - // here. - log.makeAlert("Timed out after [%d] millis waiting for handoffs", tuningConfig.getHandoffConditionTimeout()) - .addData("TaskId", task.getId()) - .emit(); - } - } - - for (SegmentsAndMetadata handedOff : handedOffList) { - log.info( - "Handoff completed for segments[%s] with metadata[%s].", - Joiner.on(", ").join( - handedOff.getSegments().stream().map(DataSegment::getIdentifier).collect(Collectors.toList()) - ), - Preconditions.checkNotNull(handedOff.getCommitMetadata(), "commitMetadata") - ); - } - - appenderator.close(); - } - catch (InterruptedException | RejectedExecutionException e) { - // (2) catch InterruptedException and RejectedExecutionException thrown for the whole ingestion steps including - // the final publishing. - caughtExceptionOuter = e; - try { - Futures.allAsList(publishWaitList).cancel(true); - Futures.allAsList(handOffWaitList).cancel(true); - if (appenderator != null) { - appenderator.closeNow(); - } - } - catch (Exception e2) { - e.addSuppressed(e2); - } - - // handle the InterruptedException that gets wrapped in a RejectedExecutionException - if (e instanceof RejectedExecutionException - && (e.getCause() == null || !(e.getCause() instanceof InterruptedException))) { - throw e; - } - - // if we were interrupted because we were asked to stop, handle the exception and return success, else rethrow - if (!stopRequested.get()) { - Thread.currentThread().interrupt(); - throw e; - } - - log.info("The task was asked to stop before completing"); - } - catch (Exception e) { - // (3) catch all other exceptions thrown for the whole ingestion steps including the final publishing. - caughtExceptionOuter = e; - try { - Futures.allAsList(publishWaitList).cancel(true); - Futures.allAsList(handOffWaitList).cancel(true); - if (appenderator != null) { - appenderator.closeNow(); - } - } - catch (Exception e2) { - e.addSuppressed(e2); - } - throw e; - } - finally { - try { - if (driver != null) { - driver.close(); - } - if (chatHandlerProvider.isPresent()) { - chatHandlerProvider.get().unregister(task.getId()); - } - - toolbox.getDruidNodeAnnouncer().unannounce(discoveryDruidNode); - toolbox.getDataSegmentServerAnnouncer().unannounce(); - } - catch (Exception e) { - if (caughtExceptionOuter != null) { - caughtExceptionOuter.addSuppressed(e); - } else { - throw e; - } - } - } - - toolbox.getTaskReportFileWriter().write(getTaskCompletionReports(null)); - return TaskStatus.success(task.getId()); - } - - private void checkPublishAndHandoffFailure() throws ExecutionException, InterruptedException - { - // Check if any publishFuture failed. - final List> publishFinished = publishWaitList - .stream() - .filter(Future::isDone) - .collect(Collectors.toList()); - - for (ListenableFuture publishFuture : publishFinished) { - // If publishFuture failed, the below line will throw an exception and catched by (1), and then (2) or (3). - publishFuture.get(); - } - - publishWaitList.removeAll(publishFinished); - - // Check if any handoffFuture failed. - final List> handoffFinished = handOffWaitList - .stream() - .filter(Future::isDone) - .collect(Collectors.toList()); - - for (ListenableFuture handoffFuture : handoffFinished) { - // If handoffFuture failed, the below line will throw an exception and catched by (1), and then (2) or (3). - handoffFuture.get(); - } - - handOffWaitList.removeAll(handoffFinished); - } - - private void publishAndRegisterHandoff(SequenceMetadata sequenceMetadata) - { - log.info("Publishing segments for sequence [%s]", sequenceMetadata); - - final ListenableFuture publishFuture = Futures.transform( - driver.publish( - sequenceMetadata.createPublisher(toolbox, ioConfig.isUseTransaction()), - sequenceMetadata.getCommitterSupplier(topic, lastPersistedOffsets).get(), - Collections.singletonList(sequenceMetadata.getSequenceName()) - ), - (Function) publishedSegmentsAndMetadata -> { - if (publishedSegmentsAndMetadata == null) { - throw new ISE( - "Transaction failure publishing segments for sequence [%s]", - sequenceMetadata - ); - } else { - return publishedSegmentsAndMetadata; - } - } - ); - publishWaitList.add(publishFuture); - - // Create a handoffFuture for every publishFuture. The created handoffFuture must fail if publishFuture fails. - final SettableFuture handoffFuture = SettableFuture.create(); - handOffWaitList.add(handoffFuture); - - Futures.addCallback( - publishFuture, - new FutureCallback() - { - @Override - public void onSuccess(SegmentsAndMetadata publishedSegmentsAndMetadata) - { - log.info( - "Published segments[%s] with metadata[%s].", - publishedSegmentsAndMetadata.getSegments() - .stream() - .map(DataSegment::getIdentifier) - .collect(Collectors.toList()), - Preconditions.checkNotNull(publishedSegmentsAndMetadata.getCommitMetadata(), "commitMetadata") - ); - - sequences.remove(sequenceMetadata); - publishingSequences.remove(sequenceMetadata.getSequenceName()); - try { - persistSequences(); - } - catch (IOException e) { - log.error(e, "Unable to persist state, dying"); - handoffFuture.setException(e); - throw new RuntimeException(e); - } - - Futures.transform( - driver.registerHandoff(publishedSegmentsAndMetadata), - new Function() - { - @Nullable - @Override - public Void apply(@Nullable SegmentsAndMetadata handoffSegmentsAndMetadata) - { - if (handoffSegmentsAndMetadata == null) { - log.warn( - "Failed to handoff segments[%s]", - publishedSegmentsAndMetadata.getSegments() - .stream() - .map(DataSegment::getIdentifier) - .collect(Collectors.toList()) - ); - } - handoffFuture.set(handoffSegmentsAndMetadata); - return null; - } - } - ); - } - - @Override - public void onFailure(Throwable t) - { - log.error(t, "Error while publishing segments for sequence[%s]", sequenceMetadata); - handoffFuture.setException(t); - } - } - ); - } - - private static File getSequencesPersistFile(TaskToolbox toolbox) - { - return new File(toolbox.getPersistDir(), "sequences.json"); - } - - private boolean restoreSequences() throws IOException - { - final File sequencesPersistFile = getSequencesPersistFile(toolbox); - if (sequencesPersistFile.exists()) { - sequences = new CopyOnWriteArrayList<>( - toolbox.getObjectMapper().>readValue( - sequencesPersistFile, - new TypeReference>() - { - } - ) - ); - return true; - } else { - return false; - } - } - - private synchronized void persistSequences() throws IOException - { - log.info("Persisting Sequences Metadata [%s]", sequences); - toolbox.getObjectMapper().writerWithType( - new TypeReference>() - { - } - ).writeValue(getSequencesPersistFile(toolbox), sequences); - } - - private Map getTaskCompletionReports(@Nullable String errorMsg) - { - return TaskReport.buildTaskReports( - new IngestionStatsAndErrorsTaskReport( - task.getId(), - new IngestionStatsAndErrorsTaskReportData( - ingestionState, - getTaskCompletionUnparseableEvents(), - getTaskCompletionRowStats(), - errorMsg - ) - ) - ); - } - - private Map getTaskCompletionUnparseableEvents() - { - Map unparseableEventsMap = new HashMap<>(); - List buildSegmentsParseExceptionMessages = IndexTaskUtils.getMessagesFromSavedParseExceptions( - savedParseExceptions - ); - if (buildSegmentsParseExceptionMessages != null) { - unparseableEventsMap.put(RowIngestionMeters.BUILD_SEGMENTS, buildSegmentsParseExceptionMessages); - } - return unparseableEventsMap; - } - - private Map getTaskCompletionRowStats() - { - Map metrics = new HashMap<>(); - metrics.put( - RowIngestionMeters.BUILD_SEGMENTS, - rowIngestionMeters.getTotals() - ); - return metrics; - } - - private void maybePersistAndPublishSequences(Supplier committerSupplier) - throws InterruptedException - { - for (SequenceMetadata sequenceMetadata : sequences) { - sequenceMetadata.updateAssignments(nextOffsets); - if (!sequenceMetadata.isOpen() && !publishingSequences.contains(sequenceMetadata.getSequenceName())) { - publishingSequences.add(sequenceMetadata.getSequenceName()); - try { - Object result = driver.persist(committerSupplier.get()); - log.info( - "Persist completed with results: [%s], adding sequence [%s] to publish queue", - result, - sequenceMetadata - ); - publishAndRegisterHandoff(sequenceMetadata); - } - catch (InterruptedException e) { - log.warn("Interrupted while persisting sequence [%s]", sequenceMetadata); - throw e; - } - } - } - } - - private Set assignPartitionsAndSeekToNext(KafkaConsumer consumer, String topic) - { - // Initialize consumer assignment. - final Set assignment = new HashSet<>(); - for (Map.Entry entry : nextOffsets.entrySet()) { - final long endOffset = endOffsets.get(entry.getKey()); - if (entry.getValue() < endOffset) { - assignment.add(entry.getKey()); - } else if (entry.getValue() == endOffset) { - log.info("Finished reading partition[%d].", entry.getKey()); - } else { - throw new ISE( - "WTF?! Cannot start from offset[%,d] > endOffset[%,d]", - entry.getValue(), - endOffset - ); - } + @Nonnull + @Override + protected List> getRecords( + RecordSupplier recordSupplier, + TaskToolbox toolbox + ) throws Exception + { + // Handles OffsetOutOfRangeException, which is thrown if the seeked-to + // offset is not present in the topic-partition. This can happen if we're asking a task to read from data + // that has not been written yet (which is totally legitimate). So let's wait for it to show up. + List> records = new ArrayList<>(); + try { + records = recordSupplier.poll(KafkaIndexTask.POLL_TIMEOUT_MILLIS); } - - KafkaIndexTask.assignPartitions(consumer, topic, assignment); - - // Seek to starting offsets. - for (final int partition : assignment) { - final long offset = nextOffsets.get(partition); - log.info("Seeking partition[%d] to offset[%,d].", partition, offset); - consumer.seek(new TopicPartition(topic, partition), offset); + catch (OffsetOutOfRangeException e) { + log.warn("OffsetOutOfRangeException with message [%s]", e.getMessage()); + possiblyResetOffsetsOrWait(e.offsetOutOfRangePartitions(), recordSupplier, toolbox); } - return assignment; + return records; } - /** - * Checks if the pauseRequested flag was set and if so blocks until pauseRequested is cleared. - *

- * Sets paused = true and signals paused so callers can be notified when the pause command has been accepted. - *

- * - * @return true if a pause request was handled, false otherwise - */ - private boolean possiblyPause() throws InterruptedException + @Override + protected SeekableStreamPartitions deserializeSeekableStreamPartitionsFromMetadata( + ObjectMapper mapper, + Object object + ) { - pauseLock.lockInterruptibly(); - try { - if (pauseRequested) { - status = Status.PAUSED; - hasPaused.signalAll(); - - while (pauseRequested) { - log.info("Pausing ingestion until resumed"); - shouldResume.await(); - } - - status = Status.READING; - shouldResume.signalAll(); - log.info("Ingestion loop resumed"); - return true; - } - } - finally { - pauseLock.unlock(); - } - - return false; + return mapper.convertValue(object, mapper.getTypeFactory().constructParametrizedType( + SeekableStreamPartitions.class, + SeekableStreamPartitions.class, + Integer.class, + Long.class + )); } private void possiblyResetOffsetsOrWait( Map outOfRangePartitions, - KafkaConsumer consumer, + RecordSupplier recordSupplier, TaskToolbox taskToolbox ) throws InterruptedException, IOException { final Map resetPartitions = new HashMap<>(); boolean doReset = false; - if (tuningConfig.isResetOffsetAutomatically()) { + if (task.getTuningConfig().isResetOffsetAutomatically()) { for (Map.Entry outOfRangePartition : outOfRangePartitions.entrySet()) { final TopicPartition topicPartition = outOfRangePartition.getKey(); final long nextOffset = outOfRangePartition.getValue(); // seek to the beginning to get the least available offset - consumer.seekToBeginning(Collections.singletonList(topicPartition)); - final long leastAvailableOffset = consumer.position(topicPartition); + StreamPartition streamPartition = StreamPartition.of( + topicPartition.topic(), + topicPartition.partition() + ); + final Long leastAvailableOffset = recordSupplier.getEarliestSequenceNumber(streamPartition); + if (leastAvailableOffset == null) { + throw new ISE( + "got null sequence number for partition[%s] when fetching from kafka!", + topicPartition.partition() + ); + } // reset the seek - consumer.seek(topicPartition, nextOffset); + recordSupplier.seek(streamPartition, nextOffset); // Reset consumer offset if resetOffsetAutomatically is set to true // and the current message offset in the kafka partition is more than the // next message offset that we are trying to fetch @@ -1067,7 +161,12 @@ private void possiblyResetOffsetsOrWait( } if (doReset) { - sendResetRequestAndWait(resetPartitions, taskToolbox); + sendResetRequestAndWait(resetPartitions.entrySet() + .stream() + .collect(Collectors.toMap(x -> StreamPartition.of( + x.getKey().topic(), + x.getKey().partition() + ), Map.Entry::getValue)), taskToolbox); } else { log.warn("Retrying in %dms", task.getPollRetryMs()); pollRetryLock.lockInterruptibly(); @@ -1083,721 +182,56 @@ private void possiblyResetOffsetsOrWait( } } - private void handleParseException(ParseException pe, ConsumerRecord record) - { - if (pe.isFromPartiallyValidRow()) { - rowIngestionMeters.incrementProcessedWithError(); - } else { - rowIngestionMeters.incrementUnparseable(); - } - - if (tuningConfig.isLogParseExceptions()) { - log.error( - pe, - "Encountered parse exception on row from partition[%d] offset[%d]", - record.partition(), - record.offset() - ); - } - - if (savedParseExceptions != null) { - savedParseExceptions.add(pe); - } - - if (rowIngestionMeters.getUnparseable() + rowIngestionMeters.getProcessedWithError() - > tuningConfig.getMaxParseExceptions()) { - log.error("Max parse exceptions exceeded, terminating task..."); - throw new RuntimeException("Max parse exceptions exceeded, terminating task..."); - } - } - - private boolean isPaused() - { - return status == Status.PAUSED; - } - - private void requestPause() - { - pauseRequested = true; - } - - private void sendResetRequestAndWait(Map outOfRangePartitions, TaskToolbox taskToolbox) - throws IOException - { - Map partitionOffsetMap = new HashMap<>(); - for (Map.Entry outOfRangePartition : outOfRangePartitions.entrySet()) { - partitionOffsetMap.put(outOfRangePartition.getKey().partition(), outOfRangePartition.getValue()); - } - boolean result = taskToolbox.getTaskActionClient() - .submit(new ResetDataSourceMetadataAction( - task.getDataSource(), - new KafkaDataSourceMetadata(new KafkaPartitions( - ioConfig.getStartPartitions() - .getTopic(), - partitionOffsetMap - )) - )); - - if (result) { - log.makeAlert("Resetting Kafka offsets for datasource [%s]", task.getDataSource()) - .addData("partitions", partitionOffsetMap.keySet()) - .emit(); - // wait for being killed by supervisor - requestPause(); - } else { - log.makeAlert("Failed to send reset request for partitions [%s]", partitionOffsetMap.keySet()).emit(); - } - } - - /** - * Authorizes action to be performed on this task's datasource - * - * @return authorization result - */ - private Access authorizationCheck(final HttpServletRequest req, Action action) - { - return IndexTaskUtils.datasourceAuthorizationCheck(req, action, task.getDataSource(), authorizerMapper); - } - @Override - public Appenderator getAppenderator() - { - return appenderator; - } - - @Override - public RowIngestionMeters getRowIngestionMeters() - { - return rowIngestionMeters; - } - - @Override - public void stopGracefully() - { - log.info("Stopping gracefully (status: [%s])", status); - stopRequested.set(true); - - synchronized (statusLock) { - if (status == Status.PUBLISHING) { - runThread.interrupt(); - return; - } - } - - try { - if (pauseLock.tryLock(KafkaIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { - try { - if (pauseRequested) { - pauseRequested = false; - shouldResume.signalAll(); - } - } - finally { - pauseLock.unlock(); - } - } else { - log.warn("While stopping: failed to acquire pauseLock before timeout, interrupting run thread"); - runThread.interrupt(); - return; - } - - if (pollRetryLock.tryLock(KafkaIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { - try { - isAwaitingRetry.signalAll(); - } - finally { - pollRetryLock.unlock(); - } - } else { - log.warn("While stopping: failed to acquire pollRetryLock before timeout, interrupting run thread"); - runThread.interrupt(); - } - } - catch (Exception e) { - Throwables.propagate(e); - } - } - - @POST - @Path("/stop") - public Response stop(@Context final HttpServletRequest req) - { - authorizationCheck(req, Action.WRITE); - stopGracefully(); - return Response.status(Response.Status.OK).build(); - } - - @GET - @Path("/status") - @Produces(MediaType.APPLICATION_JSON) - public Status getStatusHTTP(@Context final HttpServletRequest req) - { - authorizationCheck(req, Action.READ); - return status; - } - - @Override - public Status getStatus() - { - return status; - } - - @GET - @Path("/offsets/current") - @Produces(MediaType.APPLICATION_JSON) - public Map getCurrentOffsets(@Context final HttpServletRequest req) + protected SeekableStreamDataSourceMetadata createDataSourceMetadata( + SeekableStreamPartitions partitions + ) { - authorizationCheck(req, Action.READ); - return getCurrentOffsets(); + return new KafkaDataSourceMetadata(partitions); } @Override - public Map getCurrentOffsets() - { - return nextOffsets; - } - - @GET - @Path("/offsets/end") - @Produces(MediaType.APPLICATION_JSON) - public Map getEndOffsetsHTTP(@Context final HttpServletRequest req) + protected OrderedSequenceNumber createSequenceNumber(Long sequenceNumber) { - authorizationCheck(req, Action.READ); - return getEndOffsets(); + return KafkaSequenceNumber.of(sequenceNumber); } @Override - public Map getEndOffsets() - { - return endOffsets; - } - - @POST - @Path("/offsets/end") - @Consumes(MediaType.APPLICATION_JSON) - @Produces(MediaType.APPLICATION_JSON) - public Response setEndOffsetsHTTP( - Map offsets, - @QueryParam("finish") @DefaultValue("true") final boolean finish, - // this field is only for internal purposes, shouldn't be usually set by users - @Context final HttpServletRequest req - ) throws InterruptedException - { - authorizationCheck(req, Action.WRITE); - return setEndOffsets(offsets, finish); - } - - @GET - @Path("/rowStats") - @Produces(MediaType.APPLICATION_JSON) - public Response getRowStats( - @Context final HttpServletRequest req - ) - { - authorizationCheck(req, Action.READ); - Map returnMap = new HashMap<>(); - Map totalsMap = new HashMap<>(); - Map averagesMap = new HashMap<>(); - - totalsMap.put( - RowIngestionMeters.BUILD_SEGMENTS, - rowIngestionMeters.getTotals() - ); - averagesMap.put( - RowIngestionMeters.BUILD_SEGMENTS, - rowIngestionMeters.getMovingAverages() - ); - - returnMap.put("movingAverages", averagesMap); - returnMap.put("totals", totalsMap); - return Response.ok(returnMap).build(); - } - - @GET - @Path("/unparseableEvents") - @Produces(MediaType.APPLICATION_JSON) - public Response getUnparseableEvents( - @Context final HttpServletRequest req + protected void possiblyResetDataSourceMetadata( + TaskToolbox toolbox, + RecordSupplier recordSupplier, + Set> assignment, + Map currOffsets ) { - authorizationCheck(req, Action.READ); - List events = IndexTaskUtils.getMessagesFromSavedParseExceptions(savedParseExceptions); - return Response.ok(events).build(); + // do nothing } @Override - public Response setEndOffsets( - Map offsets, - final boolean finish // this field is only for internal purposes, shouldn't be usually set by users - ) throws InterruptedException - { - if (offsets == null) { - return Response.status(Response.Status.BAD_REQUEST) - .entity("Request body must contain a map of { partition:endOffset }") - .build(); - } else if (!endOffsets.keySet().containsAll(offsets.keySet())) { - return Response.status(Response.Status.BAD_REQUEST) - .entity( - StringUtils.format( - "Request contains partitions not being handled by this task, my partitions: %s", - endOffsets.keySet() - ) - ) - .build(); - } else { - try { - pauseLock.lockInterruptibly(); - // Perform all sequence related checks before checking for isPaused() - // and after acquiring pauseLock to correctly guard against duplicate requests - Preconditions.checkState(sequences.size() > 0, "WTH?! No Sequences found to set end offsets"); - - final SequenceMetadata latestSequence = sequences.get(sequences.size() - 1); - if ((latestSequence.getStartOffsets().equals(offsets) && !finish) || - (latestSequence.getEndOffsets().equals(offsets) && finish)) { - log.warn("Ignoring duplicate request, end offsets already set for sequences [%s]", sequences); - return Response.ok(offsets).build(); - } else if (latestSequence.isCheckpointed()) { - return Response.status(Response.Status.BAD_REQUEST) - .entity(StringUtils.format( - "WTH?! Sequence [%s] has already endOffsets set, cannot set to [%s]", - latestSequence, - offsets - )).build(); - } else if (!isPaused()) { - return Response.status(Response.Status.BAD_REQUEST) - .entity("Task must be paused before changing the end offsets") - .build(); - } - - for (Map.Entry entry : offsets.entrySet()) { - if (entry.getValue().compareTo(nextOffsets.get(entry.getKey())) < 0) { - return Response.status(Response.Status.BAD_REQUEST) - .entity( - StringUtils.format( - "End offset must be >= current offset for partition [%s] (current: %s)", - entry.getKey(), - nextOffsets.get(entry.getKey()) - ) - ) - .build(); - } - } - - resetNextCheckpointTime(); - latestSequence.setEndOffsets(offsets); - - if (finish) { - log.info("Updating endOffsets from [%s] to [%s]", endOffsets, offsets); - endOffsets.putAll(offsets); - } else { - // create new sequence - final SequenceMetadata newSequence = new SequenceMetadata( - latestSequence.getSequenceId() + 1, - StringUtils.format("%s_%d", ioConfig.getBaseSequenceName(), latestSequence.getSequenceId() + 1), - offsets, - endOffsets, - false - ); - sequences.add(newSequence); - } - - persistSequences(); - } - catch (Exception e) { - log.error(e, "Unable to set end offsets, dying"); - backgroundThreadException = e; - // should resume to immediately finish kafka index task as failed - resume(); - return Response.status(Response.Status.INTERNAL_SERVER_ERROR) - .entity(Throwables.getStackTraceAsString(e)) - .build(); - } - finally { - pauseLock.unlock(); - } - } - - resume(); - - return Response.ok(offsets).build(); - } - - private void resetNextCheckpointTime() - { - nextCheckpointTime = DateTimes.nowUtc().plus(tuningConfig.getIntermediateHandoffPeriod()).getMillis(); - } - - @GET - @Path("/checkpoints") - @Produces(MediaType.APPLICATION_JSON) - public Map> getCheckpointsHTTP(@Context final HttpServletRequest req) + protected boolean isEndSequenceOffsetsExclusive() { - authorizationCheck(req, Action.READ); - return getCheckpoints(); - } - - private Map> getCheckpoints() - { - TreeMap> result = new TreeMap<>(); - result.putAll( - sequences.stream().collect(Collectors.toMap(SequenceMetadata::getSequenceId, SequenceMetadata::getStartOffsets)) - ); - return result; - } - - /** - * Signals the ingestion loop to pause. - * - * @return one of the following Responses: 400 Bad Request if the task has started publishing; 202 Accepted if the - * method has timed out and returned before the task has paused; 200 OK with a map of the current partition offsets - * in the response body if the task successfully paused - */ - @POST - @Path("/pause") - @Produces(MediaType.APPLICATION_JSON) - public Response pauseHTTP( - @Context final HttpServletRequest req - ) throws InterruptedException - { - authorizationCheck(req, Action.WRITE); - return pause(); + return true; } @Override - public Response pause() throws InterruptedException - { - if (!(status == Status.PAUSED || status == Status.READING)) { - return Response.status(Response.Status.BAD_REQUEST) - .entity(StringUtils.format("Can't pause, task is not in a pausable state (state: [%s])", status)) - .build(); - } - - pauseLock.lockInterruptibly(); - try { - pauseRequested = true; - - pollRetryLock.lockInterruptibly(); - try { - isAwaitingRetry.signalAll(); - } - finally { - pollRetryLock.unlock(); - } - - if (isPaused()) { - shouldResume.signalAll(); // kick the monitor so it re-awaits with the new pauseMillis - } - - long nanos = TimeUnit.SECONDS.toNanos(2); - while (!isPaused()) { - if (nanos <= 0L) { - return Response.status(Response.Status.ACCEPTED) - .entity("Request accepted but task has not yet paused") - .build(); - } - nanos = hasPaused.awaitNanos(nanos); - } - } - finally { - pauseLock.unlock(); - } - - try { - return Response.ok().entity(toolbox.getObjectMapper().writeValueAsString(getCurrentOffsets())).build(); - } - catch (JsonProcessingException e) { - throw Throwables.propagate(e); - } - } - - @POST - @Path("/resume") - public Response resumeHTTP(@Context final HttpServletRequest req) throws InterruptedException + protected boolean isStartingSequenceOffsetsExclusive() { - authorizationCheck(req, Action.WRITE); - resume(); - return Response.status(Response.Status.OK).build(); + return false; } @Override - public void resume() throws InterruptedException - { - pauseLock.lockInterruptibly(); - try { - pauseRequested = false; - shouldResume.signalAll(); - - long nanos = TimeUnit.SECONDS.toNanos(5); - while (isPaused()) { - if (nanos <= 0L) { - throw new RuntimeException("Resume command was not accepted within 5 seconds"); - } - nanos = shouldResume.awaitNanos(nanos); - } - } - finally { - pauseLock.unlock(); - } - } - - @GET - @Path("/time/start") - @Produces(MediaType.APPLICATION_JSON) - public DateTime getStartTime(@Context final HttpServletRequest req) - { - authorizationCheck(req, Action.WRITE); - return startTime; - } - - private static class SequenceMetadata + protected boolean isEndOfShard(Long seqNum) { - /** - * Lock for accessing {@link #endOffsets} and {@link #checkpointed}. This lock is required because - * {@link #setEndOffsets)} can be called by both the main thread and the HTTP thread. - */ - private final ReentrantLock lock = new ReentrantLock(); - - private final int sequenceId; - private final String sequenceName; - private final Map startOffsets; - private final Map endOffsets; - private final Set assignments; - private final boolean sentinel; - private boolean checkpointed; - - @JsonCreator - public SequenceMetadata( - @JsonProperty("sequenceId") int sequenceId, - @JsonProperty("sequenceName") String sequenceName, - @JsonProperty("startOffsets") Map startOffsets, - @JsonProperty("endOffsets") Map endOffsets, - @JsonProperty("checkpointed") boolean checkpointed - ) - { - Preconditions.checkNotNull(sequenceName); - Preconditions.checkNotNull(startOffsets); - Preconditions.checkNotNull(endOffsets); - this.sequenceId = sequenceId; - this.sequenceName = sequenceName; - this.startOffsets = ImmutableMap.copyOf(startOffsets); - this.endOffsets = new HashMap<>(endOffsets); - this.assignments = new HashSet<>(startOffsets.keySet()); - this.checkpointed = checkpointed; - this.sentinel = false; - } - - @JsonProperty - public int getSequenceId() - { - return sequenceId; - } - - @JsonProperty - public boolean isCheckpointed() - { - lock.lock(); - try { - return checkpointed; - } - finally { - lock.unlock(); - } - } - - @JsonProperty - public String getSequenceName() - { - return sequenceName; - } - - @JsonProperty - public Map getStartOffsets() - { - return startOffsets; - } - - @JsonProperty - public Map getEndOffsets() - { - lock.lock(); - try { - return endOffsets; - } - finally { - lock.unlock(); - } - } - - @JsonProperty - public boolean isSentinel() - { - return sentinel; - } - - void setEndOffsets(Map newEndOffsets) - { - lock.lock(); - try { - endOffsets.putAll(newEndOffsets); - checkpointed = true; - } - finally { - lock.unlock(); - } - } - - void updateAssignments(Map nextPartitionOffset) - { - lock.lock(); - try { - assignments.clear(); - nextPartitionOffset.forEach((key, value) -> { - if (Longs.compare(endOffsets.get(key), nextPartitionOffset.get(key)) > 0) { - assignments.add(key); - } - }); - } - finally { - lock.unlock(); - } - } - - boolean isOpen() - { - return !assignments.isEmpty(); - } - - boolean canHandle(ConsumerRecord record) - { - lock.lock(); - try { - final Long partitionEndOffset = endOffsets.get(record.partition()); - return isOpen() - && partitionEndOffset != null - && record.offset() >= startOffsets.get(record.partition()) - && record.offset() < partitionEndOffset; - } - finally { - lock.unlock(); - } - } - - @Override - public String toString() - { - lock.lock(); - try { - return "SequenceMetadata{" + - "sequenceName='" + sequenceName + '\'' + - ", sequenceId=" + sequenceId + - ", startOffsets=" + startOffsets + - ", endOffsets=" + endOffsets + - ", assignments=" + assignments + - ", sentinel=" + sentinel + - ", checkpointed=" + checkpointed + - '}'; - } - finally { - lock.unlock(); - } - } - - Supplier getCommitterSupplier(String topic, Map lastPersistedOffsets) - { - // Set up committer. - return () -> - new Committer() - { - @Override - public Object getMetadata() - { - lock.lock(); - - try { - Preconditions.checkState( - assignments.isEmpty(), - "This committer can be used only once all the records till offsets [%s] have been consumed, also make" - + " sure to call updateAssignments before using this committer", - endOffsets - ); - - // merge endOffsets for this sequence with globally lastPersistedOffsets - // This is done because this committer would be persisting only sub set of segments - // corresponding to the current sequence. Generally, lastPersistedOffsets should already - // cover endOffsets but just to be sure take max of offsets and persist that - for (Map.Entry partitionOffset : endOffsets.entrySet()) { - lastPersistedOffsets.put( - partitionOffset.getKey(), - Math.max( - partitionOffset.getValue(), - lastPersistedOffsets.getOrDefault(partitionOffset.getKey(), 0L) - ) - ); - } - - // Publish metadata can be different from persist metadata as we are going to publish only - // subset of segments - return ImmutableMap.of( - METADATA_NEXT_PARTITIONS, new KafkaPartitions(topic, lastPersistedOffsets), - METADATA_PUBLISH_PARTITIONS, new KafkaPartitions(topic, endOffsets) - ); - } - finally { - lock.unlock(); - } - } - - @Override - public void run() - { - // Do nothing. - } - }; - } - - TransactionalSegmentPublisher createPublisher(TaskToolbox toolbox, boolean useTransaction) - { - return (segments, commitMetadata) -> { - final KafkaPartitions finalPartitions = toolbox.getObjectMapper().convertValue( - ((Map) Preconditions.checkNotNull(commitMetadata, "commitMetadata")).get(METADATA_PUBLISH_PARTITIONS), - KafkaPartitions.class - ); - - // Sanity check, we should only be publishing things that match our desired end state. - if (!getEndOffsets().equals(finalPartitions.getPartitionOffsetMap())) { - throw new ISE( - "WTF?! Driver for sequence [%s], attempted to publish invalid metadata[%s].", - toString(), - commitMetadata - ); - } - - final SegmentTransactionalInsertAction action; - - if (useTransaction) { - action = new SegmentTransactionalInsertAction( - segments, - new KafkaDataSourceMetadata(new KafkaPartitions(finalPartitions.getTopic(), getStartOffsets())), - new KafkaDataSourceMetadata(finalPartitions) - ); - } else { - action = new SegmentTransactionalInsertAction(segments, null, null); - } - - log.info("Publishing with isTransaction[%s].", useTransaction); - - return toolbox.getTaskActionClient().submit(action); - }; - } + return false; } @Nullable - private static TreeMap> getCheckPointsFromContext( + @Override + protected TreeMap> getCheckPointsFromContext( TaskToolbox toolbox, - KafkaIndexTask task + String checkpointsString ) throws IOException { - final String checkpointsString = task.getContextValue("checkpoints"); if (checkpointsString != null) { log.info("Checkpoints [%s]", checkpointsString); return toolbox.getObjectMapper().readValue( @@ -1811,3 +245,4 @@ private static TreeMap> getCheckPointsFromContext( } } } + diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadata.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadata.java index 16909892771a..55b03c000f30 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadata.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadata.java @@ -21,133 +21,28 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import org.apache.druid.indexing.overlord.DataSourceMetadata; -import org.apache.druid.java.util.common.IAE; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; -import java.util.HashMap; import java.util.Map; -import java.util.Objects; -public class KafkaDataSourceMetadata implements DataSourceMetadata +public class KafkaDataSourceMetadata extends SeekableStreamDataSourceMetadata { - private final KafkaPartitions kafkaPartitions; @JsonCreator public KafkaDataSourceMetadata( - @JsonProperty("partitions") KafkaPartitions kafkaPartitions + @JsonProperty("partitions") SeekableStreamPartitions kafkaPartitions ) { - this.kafkaPartitions = kafkaPartitions; - } - - @JsonProperty("partitions") - public KafkaPartitions getKafkaPartitions() - { - return kafkaPartitions; - } - - @Override - public boolean isValidStart() - { - return true; - } - - @Override - public boolean matches(DataSourceMetadata other) - { - if (getClass() != other.getClass()) { - return false; - } - - return plus(other).equals(other.plus(this)); - } - - @Override - public DataSourceMetadata plus(DataSourceMetadata other) - { - if (!(other instanceof KafkaDataSourceMetadata)) { - throw new IAE( - "Expected instance of %s, got %s", - KafkaDataSourceMetadata.class.getCanonicalName(), - other.getClass().getCanonicalName() - ); - } - - final KafkaDataSourceMetadata that = (KafkaDataSourceMetadata) other; - - if (that.getKafkaPartitions().getTopic().equals(kafkaPartitions.getTopic())) { - // Same topic, merge offsets. - final Map newMap = new HashMap<>(); - - for (Map.Entry entry : kafkaPartitions.getPartitionOffsetMap().entrySet()) { - newMap.put(entry.getKey(), entry.getValue()); - } - - for (Map.Entry entry : that.getKafkaPartitions().getPartitionOffsetMap().entrySet()) { - newMap.put(entry.getKey(), entry.getValue()); - } - - return new KafkaDataSourceMetadata(new KafkaPartitions(kafkaPartitions.getTopic(), newMap)); - } else { - // Different topic, prefer "other". - return other; - } - } - - @Override - public DataSourceMetadata minus(DataSourceMetadata other) - { - if (!(other instanceof KafkaDataSourceMetadata)) { - throw new IAE( - "Expected instance of %s, got %s", - KafkaDataSourceMetadata.class.getCanonicalName(), - other.getClass().getCanonicalName() - ); - } - - final KafkaDataSourceMetadata that = (KafkaDataSourceMetadata) other; - - if (that.getKafkaPartitions().getTopic().equals(kafkaPartitions.getTopic())) { - // Same topic, remove partitions present in "that" from "this" - final Map newMap = new HashMap<>(); - - for (Map.Entry entry : kafkaPartitions.getPartitionOffsetMap().entrySet()) { - if (!that.getKafkaPartitions().getPartitionOffsetMap().containsKey(entry.getKey())) { - newMap.put(entry.getKey(), entry.getValue()); - } - } - - return new KafkaDataSourceMetadata(new KafkaPartitions(kafkaPartitions.getTopic(), newMap)); - } else { - // Different topic, prefer "this". - return this; - } + super(kafkaPartitions); } @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - KafkaDataSourceMetadata that = (KafkaDataSourceMetadata) o; - return Objects.equals(kafkaPartitions, that.kafkaPartitions); - } - - @Override - public int hashCode() - { - return Objects.hash(kafkaPartitions); - } - - @Override - public String toString() + protected SeekableStreamDataSourceMetadata createConcreteDataSourceMetaData( + String streamId, + Map newMap + ) { - return "KafkaDataSourceMetadata{" + - "kafkaPartitions=" + kafkaPartitions + - '}'; + return new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(streamId, newMap)); } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java index bb73651e6e87..22ede084b3a9 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java @@ -24,77 +24,32 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Joiner; -import com.google.common.base.Optional; -import com.google.common.base.Preconditions; -import com.google.common.collect.Lists; -import org.apache.druid.data.input.InputRow; -import org.apache.druid.data.input.impl.InputRowParser; -import org.apache.druid.indexer.TaskStatus; -import org.apache.druid.indexing.appenderator.ActionBasedSegmentAllocator; -import org.apache.druid.indexing.appenderator.ActionBasedUsedSegmentChecker; -import org.apache.druid.indexing.common.TaskToolbox; -import org.apache.druid.indexing.common.actions.SegmentAllocateAction; -import org.apache.druid.indexing.common.actions.TaskActionClient; import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; -import org.apache.druid.indexing.common.task.AbstractTask; -import org.apache.druid.indexing.common.task.RealtimeIndexTask; import org.apache.druid.indexing.common.task.TaskResource; -import org.apache.druid.indexing.common.task.Tasks; -import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisor; -import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorIOConfig; -import org.apache.druid.java.util.common.Intervals; -import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.parsers.ParseException; -import org.apache.druid.java.util.emitter.EmittingLogger; -import org.apache.druid.metadata.PasswordProvider; -import org.apache.druid.query.NoopQueryRunner; -import org.apache.druid.query.Query; -import org.apache.druid.query.QueryRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; import org.apache.druid.segment.indexing.DataSchema; -import org.apache.druid.segment.realtime.FireDepartmentMetrics; -import org.apache.druid.segment.realtime.appenderator.Appenderator; -import org.apache.druid.segment.realtime.appenderator.Appenderators; -import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; -import org.apache.druid.segment.realtime.firehose.ChatHandler; import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.utils.CircularBuffer; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.ByteArrayDeserializer; -import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.HashMap; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; -public class KafkaIndexTask extends AbstractTask implements ChatHandler +public class KafkaIndexTask extends SeekableStreamIndexTask { - public enum Status - { - NOT_STARTED, - STARTING, - READING, - PAUSED, - PUBLISHING - // ideally this should be called FINISHING now as the task does incremental publishes - // through out its lifetime - } - - private static final EmittingLogger log = new EmittingLogger(KafkaIndexTask.class); private static final String TYPE = "index_kafka"; static final long POLL_TIMEOUT_MILLIS = TimeUnit.MILLISECONDS.toMillis(100); - static final long LOCK_ACQUIRE_TIMEOUT_SECONDS = 15; - private final DataSchema dataSchema; - private final InputRowParser parser; - private final KafkaTuningConfig tuningConfig; - private final KafkaIOConfig ioConfig; - private final Optional chatHandlerProvider; - private final KafkaIndexTaskRunner runner; + private final KafkaIndexTaskIOConfig ioConfig; private final ObjectMapper configMapper; // This value can be tuned in some tests @@ -105,8 +60,8 @@ public KafkaIndexTask( @JsonProperty("id") String id, @JsonProperty("resource") TaskResource taskResource, @JsonProperty("dataSchema") DataSchema dataSchema, - @JsonProperty("tuningConfig") KafkaTuningConfig tuningConfig, - @JsonProperty("ioConfig") KafkaIOConfig ioConfig, + @JsonProperty("tuningConfig") KafkaIndexTaskTuningConfig tuningConfig, + @JsonProperty("ioConfig") KafkaIndexTaskIOConfig ioConfig, @JsonProperty("context") Map context, @JacksonInject ChatHandlerProvider chatHandlerProvider, @JacksonInject AuthorizerMapper authorizerMapper, @@ -115,46 +70,20 @@ public KafkaIndexTask( ) { super( - id == null ? makeTaskId(dataSchema.getDataSource()) : id, - StringUtils.format("%s_%s", TYPE, dataSchema.getDataSource()), + id == null ? getFormattedId(dataSchema.getDataSource(), TYPE) : id, taskResource, - dataSchema.getDataSource(), - context + dataSchema, + tuningConfig, + ioConfig, + context, + chatHandlerProvider, + authorizerMapper, + rowIngestionMetersFactory, + getFormattedGroupId(dataSchema.getDataSource(), TYPE) ); - - this.dataSchema = Preconditions.checkNotNull(dataSchema, "dataSchema"); - this.parser = Preconditions.checkNotNull((InputRowParser) dataSchema.getParser(), "parser"); - this.tuningConfig = Preconditions.checkNotNull(tuningConfig, "tuningConfig"); - this.ioConfig = Preconditions.checkNotNull(ioConfig, "ioConfig"); - this.chatHandlerProvider = Optional.fromNullable(chatHandlerProvider); this.configMapper = configMapper; - final CircularBuffer savedParseExceptions; - if (tuningConfig.getMaxSavedParseExceptions() > 0) { - savedParseExceptions = new CircularBuffer<>(tuningConfig.getMaxSavedParseExceptions()); - } else { - savedParseExceptions = null; - } + this.ioConfig = ioConfig; - if (context != null && context.get(KafkaSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED) != null - && ((boolean) context.get(KafkaSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED))) { - runner = new IncrementalPublishingKafkaIndexTaskRunner( - this, - parser, - authorizerMapper, - this.chatHandlerProvider, - savedParseExceptions, - rowIngestionMetersFactory - ); - } else { - runner = new LegacyKafkaIndexTaskRunner( - this, - parser, - authorizerMapper, - this.chatHandlerProvider, - savedParseExceptions, - rowIngestionMetersFactory - ); - } } long getPollRetryMs() @@ -162,126 +91,7 @@ long getPollRetryMs() return pollRetryMs; } - private static String makeTaskId(String dataSource) - { - return Joiner.on("_").join(TYPE, dataSource, RealtimeIndexTask.makeRandomId()); - } - - @Override - public int getPriority() - { - return getContextValue(Tasks.PRIORITY_KEY, Tasks.DEFAULT_REALTIME_TASK_PRIORITY); - } - - @Override - public String getType() - { - return TYPE; - } - - @Override - public boolean isReady(TaskActionClient taskActionClient) - { - return true; - } - - @JsonProperty - public DataSchema getDataSchema() - { - return dataSchema; - } - - @JsonProperty - public KafkaTuningConfig getTuningConfig() - { - return tuningConfig; - } - - @JsonProperty("ioConfig") - public KafkaIOConfig getIOConfig() - { - return ioConfig; - } - - - @Override - public TaskStatus run(final TaskToolbox toolbox) - { - return runner.run(toolbox); - } - - @Override - public boolean canRestore() - { - return true; - } - - @Override - public void stopGracefully() - { - runner.stopGracefully(); - } - - @Override - public QueryRunner getQueryRunner(Query query) - { - if (runner.getAppenderator() == null) { - // Not yet initialized, no data yet, just return a noop runner. - return new NoopQueryRunner<>(); - } - - return (queryPlus, responseContext) -> queryPlus.run(runner.getAppenderator(), responseContext); - } - - Appenderator newAppenderator(FireDepartmentMetrics metrics, TaskToolbox toolbox) - { - return Appenderators.createRealtime( - dataSchema, - tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()), - metrics, - toolbox.getSegmentPusher(), - toolbox.getObjectMapper(), - toolbox.getIndexIO(), - toolbox.getIndexMergerV9(), - toolbox.getQueryRunnerFactoryConglomerate(), - toolbox.getSegmentAnnouncer(), - toolbox.getEmitter(), - toolbox.getQueryExecutorService(), - toolbox.getCache(), - toolbox.getCacheConfig(), - toolbox.getCachePopulatorStats() - ); - } - - StreamAppenderatorDriver newDriver( - final Appenderator appenderator, - final TaskToolbox toolbox, - final FireDepartmentMetrics metrics - ) - { - return new StreamAppenderatorDriver( - appenderator, - new ActionBasedSegmentAllocator( - toolbox.getTaskActionClient(), - dataSchema, - (schema, row, sequenceName, previousSegmentId, skipSegmentLineageCheck) -> new SegmentAllocateAction( - schema.getDataSource(), - row.getTimestamp(), - schema.getGranularitySpec().getQueryGranularity(), - schema.getGranularitySpec().getSegmentGranularity(), - sequenceName, - previousSegmentId, - skipSegmentLineageCheck - ) - ), - toolbox.getSegmentHandoffNotifierFactory(), - new ActionBasedUsedSegmentChecker(toolbox.getTaskActionClient()), - toolbox.getDataSegmentKiller(), - toolbox.getObjectMapper(), - metrics - ); - } - + @Deprecated KafkaConsumer newConsumer() { ClassLoader currCtxCl = Thread.currentThread().getContextClassLoader(); @@ -290,7 +100,11 @@ KafkaConsumer newConsumer() final Properties props = new Properties(); - addConsumerPropertiesFromConfig(props, configMapper, ioConfig.getConsumerProperties()); + KafkaRecordSupplier.addConsumerPropertiesFromConfig( + props, + configMapper, + ioConfig.getConsumerProperties() + ); props.setProperty("enable.auto.commit", "false"); props.setProperty("auto.offset.reset", "none"); @@ -304,25 +118,7 @@ KafkaConsumer newConsumer() } } - public static void addConsumerPropertiesFromConfig(Properties properties, ObjectMapper configMapper, Map consumerProperties) - { - // Extract passwords before SSL connection to Kafka - for (Map.Entry entry : consumerProperties.entrySet()) { - String propertyKey = entry.getKey(); - if (propertyKey.equals(KafkaSupervisorIOConfig.TRUST_STORE_PASSWORD_KEY) - || propertyKey.equals(KafkaSupervisorIOConfig.KEY_STORE_PASSWORD_KEY) - || propertyKey.equals(KafkaSupervisorIOConfig.KEY_PASSWORD_KEY)) { - PasswordProvider configPasswordProvider = configMapper.convertValue( - entry.getValue(), - PasswordProvider.class - ); - properties.setProperty(propertyKey, configPasswordProvider.getPassword()); - } else { - properties.setProperty(propertyKey, String.valueOf(entry.getValue())); - } - } - } - + @Deprecated static void assignPartitions( final KafkaConsumer consumer, final String topic, @@ -330,45 +126,62 @@ static void assignPartitions( ) { consumer.assign( - Lists.newArrayList( + new ArrayList<>( partitions.stream().map(n -> new TopicPartition(topic, n)).collect(Collectors.toList()) ) ); } - boolean withinMinMaxRecordTime(final InputRow row) + @Override + protected SeekableStreamIndexTaskRunner createTaskRunner() { - final boolean beforeMinimumMessageTime = ioConfig.getMinimumMessageTime().isPresent() - && ioConfig.getMinimumMessageTime().get().isAfter(row.getTimestamp()); - - final boolean afterMaximumMessageTime = ioConfig.getMaximumMessageTime().isPresent() - && ioConfig.getMaximumMessageTime().get().isBefore(row.getTimestamp()); - - if (!Intervals.ETERNITY.contains(row.getTimestamp())) { - final String errorMsg = StringUtils.format( - "Encountered row with timestamp that cannot be represented as a long: [%s]", - row + if (context != null && context.get(SeekableStreamSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED) != null + && ((boolean) context.get(SeekableStreamSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED))) { + return new IncrementalPublishingKafkaIndexTaskRunner( + this, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory + ); + } else { + return new LegacyKafkaIndexTaskRunner( + this, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory ); - throw new ParseException(errorMsg); } + } + + @Override + protected KafkaRecordSupplier newTaskRecordSupplier() + { + ClassLoader currCtxCl = Thread.currentThread().getContextClassLoader(); + try { + Thread.currentThread().setContextClassLoader(getClass().getClassLoader()); + + final Map props = new HashMap<>(((KafkaIndexTaskIOConfig) super.ioConfig).getConsumerProperties()); - if (log.isDebugEnabled()) { - if (beforeMinimumMessageTime) { - log.debug( - "CurrentTimeStamp[%s] is before MinimumMessageTime[%s]", - row.getTimestamp(), - ioConfig.getMinimumMessageTime().get() - ); - } else if (afterMaximumMessageTime) { - log.debug( - "CurrentTimeStamp[%s] is after MaximumMessageTime[%s]", - row.getTimestamp(), - ioConfig.getMaximumMessageTime().get() - ); - } + props.put("auto.offset.reset", "none"); + props.put("key.deserializer", ByteArrayDeserializer.class.getName()); + props.put("value.deserializer", ByteArrayDeserializer.class.getName()); + + return new KafkaRecordSupplier(props, configMapper); + } + finally { + Thread.currentThread().setContextClassLoader(currCtxCl); } + } - return !beforeMinimumMessageTime && !afterMaximumMessageTime; + @Override + @JsonProperty + public KafkaIndexTaskTuningConfig getTuningConfig() + { + return (KafkaIndexTaskTuningConfig) super.getTuningConfig(); } @VisibleForTesting @@ -377,15 +190,16 @@ void setPollRetryMs(long retryMs) this.pollRetryMs = retryMs; } - @VisibleForTesting - Appenderator getAppenderator() + @Override + @JsonProperty("ioConfig") + public KafkaIndexTaskIOConfig getIOConfig() { - return runner.getAppenderator(); + return (KafkaIndexTaskIOConfig) super.getIOConfig(); } - @VisibleForTesting - KafkaIndexTaskRunner getRunner() + @Override + public String getType() { - return runner; + return TYPE; } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClient.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClient.java index 30f1c89a910e..a915f44e09d4 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClient.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClient.java @@ -19,37 +19,15 @@ package org.apache.druid.indexing.kafka; -import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableMap; -import com.google.common.util.concurrent.ListenableFuture; -import org.apache.druid.indexing.common.IndexTaskClient; import org.apache.druid.indexing.common.TaskInfoProvider; -import org.apache.druid.indexing.kafka.KafkaIndexTask.Status; -import org.apache.druid.java.util.common.ISE; -import org.apache.druid.java.util.common.RE; -import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.jackson.JacksonUtils; -import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClient; import org.apache.druid.java.util.http.client.HttpClient; -import org.apache.druid.java.util.http.client.response.FullResponseHolder; -import org.jboss.netty.handler.codec.http.HttpMethod; -import org.jboss.netty.handler.codec.http.HttpResponseStatus; -import org.joda.time.DateTime; import org.joda.time.Duration; -import javax.annotation.Nullable; -import java.io.IOException; -import java.util.Collections; -import java.util.Map; -import java.util.TreeMap; - -public class KafkaIndexTaskClient extends IndexTaskClient +public class KafkaIndexTaskClient extends SeekableStreamIndexTaskClient { - private static final EmittingLogger log = new EmittingLogger(KafkaIndexTaskClient.class); - private static final TreeMap> EMPTY_TREE_MAP = new TreeMap<>(); - - public KafkaIndexTaskClient( + KafkaIndexTaskClient( HttpClient httpClient, ObjectMapper jsonMapper, TaskInfoProvider taskInfoProvider, @@ -59,315 +37,26 @@ public KafkaIndexTaskClient( long numRetries ) { - super(httpClient, jsonMapper, taskInfoProvider, httpTimeout, dataSource, numThreads, numRetries); - } - - public boolean stop(final String id, final boolean publish) - { - log.debug("Stop task[%s] publish[%s]", id, publish); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent( - id, - HttpMethod.POST, - "stop", - publish ? "publish=true" : null, - true - ); - return isSuccess(response); - } - catch (NoTaskLocationException e) { - return false; - } - catch (TaskNotRunnableException e) { - log.info("Task [%s] couldn't be stopped because it is no longer running", id); - return true; - } - catch (Exception e) { - log.warn(e, "Exception while stopping task [%s]", id); - return false; - } - } - - public boolean resume(final String id) - { - log.debug("Resume task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.POST, "resume", null, true); - return isSuccess(response); - } - catch (NoTaskLocationException | IOException e) { - log.warn(e, "Exception while stopping task [%s]", id); - return false; - } - } - - public Map pause(final String id) - { - log.debug("Pause task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent( - id, - HttpMethod.POST, - "pause", - null, - true - ); - - final HttpResponseStatus responseStatus = response.getStatus(); - final String responseContent = response.getContent(); - - if (responseStatus.equals(HttpResponseStatus.OK)) { - log.info("Task [%s] paused successfully", id); - return deserialize(responseContent, new TypeReference>() - { - }); - } else if (responseStatus.equals(HttpResponseStatus.ACCEPTED)) { - // The task received the pause request, but its status hasn't been changed yet. - while (true) { - final Status status = getStatus(id); - if (status == KafkaIndexTask.Status.PAUSED) { - return getCurrentOffsets(id, true); - } - - final Duration delay = newRetryPolicy().getAndIncrementRetryDelay(); - if (delay == null) { - throw new ISE( - "Task [%s] failed to change its status from [%s] to [%s], aborting", - id, - status, - Status.PAUSED - ); - } else { - final long sleepTime = delay.getMillis(); - log.info( - "Still waiting for task [%s] to change its status to [%s]; will try again in [%s]", - id, - Status.PAUSED, - new Duration(sleepTime).toString() - ); - Thread.sleep(sleepTime); - } - } - } else { - throw new ISE( - "Pause request for task [%s] failed with response [%s] : [%s]", - id, - responseStatus, - responseContent - ); - } - } - catch (NoTaskLocationException e) { - log.error("Exception [%s] while pausing Task [%s]", e.getMessage(), id); - return ImmutableMap.of(); - } - catch (IOException | InterruptedException e) { - throw new RE(e, "Exception [%s] while pausing Task [%s]", e.getMessage(), id); - } - } - - public KafkaIndexTask.Status getStatus(final String id) - { - log.debug("GetStatus task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "status", null, true); - return deserialize(response.getContent(), KafkaIndexTask.Status.class); - } - catch (NoTaskLocationException e) { - return KafkaIndexTask.Status.NOT_STARTED; - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - @Nullable - public DateTime getStartTime(final String id) - { - log.debug("GetStartTime task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "time/start", null, true); - return response.getContent() == null || response.getContent().isEmpty() - ? null - : deserialize(response.getContent(), DateTime.class); - } - catch (NoTaskLocationException e) { - return null; - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - public Map getMovingAverages(final String id) - { - log.debug("GetMovingAverages task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent( - id, - HttpMethod.GET, - "rowStats", - null, - true - ); - return response.getContent() == null || response.getContent().isEmpty() - ? Collections.emptyMap() - : deserialize(response.getContent(), JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT); - } - catch (NoTaskLocationException e) { - return Collections.emptyMap(); - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - public Map getCurrentOffsets(final String id, final boolean retry) - { - log.debug("GetCurrentOffsets task[%s] retry[%s]", id, retry); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent( - id, - HttpMethod.GET, - "offsets/current", - null, - retry - ); - return deserialize(response.getContent(), new TypeReference>() - { - }); - } - catch (NoTaskLocationException e) { - return ImmutableMap.of(); - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - public TreeMap> getCheckpoints(final String id, final boolean retry) - { - log.debug("GetCheckpoints task[%s] retry[%s]", id, retry); - try { - final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "checkpoints", null, retry); - return deserialize( - response.getContent(), - new TypeReference>>() - { - } - ); - } - catch (NoTaskLocationException e) { - return EMPTY_TREE_MAP; - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - public ListenableFuture>> getCheckpointsAsync( - final String id, - final boolean retry - ) - { - return doAsync(() -> getCheckpoints(id, retry)); - } - - public Map getEndOffsets(final String id) - { - log.debug("GetEndOffsets task[%s]", id); - - try { - final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "offsets/end", null, true); - return deserialize(response.getContent(), new TypeReference>() - { - }); - } - catch (NoTaskLocationException e) { - return ImmutableMap.of(); - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - public boolean setEndOffsets( - final String id, - final Map endOffsets, - final boolean finalize - ) throws IOException - { - log.debug("SetEndOffsets task[%s] endOffsets[%s] finalize[%s]", id, endOffsets, finalize); - - try { - final FullResponseHolder response = submitJsonRequest( - id, - HttpMethod.POST, - "offsets/end", - StringUtils.format("finish=%s", finalize), - serialize(endOffsets), - true - ); - return isSuccess(response); - } - catch (NoTaskLocationException e) { - return false; - } + super( + httpClient, + jsonMapper, + taskInfoProvider, + dataSource, + numThreads, + httpTimeout, + numRetries + ); } - public ListenableFuture stopAsync(final String id, final boolean publish) - { - return doAsync(() -> stop(id, publish)); - } - - public ListenableFuture resumeAsync(final String id) - { - return doAsync(() -> resume(id)); - } - - public ListenableFuture> pauseAsync(final String id) - { - return doAsync(() -> pause(id)); - } - - public ListenableFuture getStatusAsync(final String id) - { - return doAsync(() -> getStatus(id)); - } - - public ListenableFuture getStartTimeAsync(final String id) - { - return doAsync(() -> getStartTime(id)); - } - - public ListenableFuture> getCurrentOffsetsAsync(final String id, final boolean retry) - { - return doAsync(() -> getCurrentOffsets(id, retry)); - } - - public ListenableFuture> getEndOffsetsAsync(final String id) - { - return doAsync(() -> getEndOffsets(id)); - } - - public ListenableFuture setEndOffsetsAsync( - final String id, - final Map endOffsets, - final boolean finalize - ) + @Override + protected Class getPartitionType() { - return doAsync(() -> setEndOffsets(id, endOffsets, finalize)); + return Integer.class; } - public ListenableFuture> getMovingAveragesAsync(final String id) + @Override + protected Class getSequenceType() { - return doAsync(() -> getMovingAverages(id)); + return Long.class; } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java index d7f7f61be4c9..cfa7c7236553 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java @@ -24,23 +24,19 @@ import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.indexing.common.TaskInfoProvider; -import org.apache.druid.indexing.common.task.IndexTaskClientFactory; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; import org.apache.druid.java.util.http.client.HttpClient; import org.joda.time.Duration; -public class KafkaIndexTaskClientFactory implements IndexTaskClientFactory +public class KafkaIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory { - private HttpClient httpClient; - private ObjectMapper mapper; - @Inject public KafkaIndexTaskClientFactory( @EscalatedGlobal HttpClient httpClient, @Json ObjectMapper mapper ) { - this.httpClient = httpClient; - this.mapper = mapper; + super(httpClient, mapper); } @Override @@ -53,8 +49,8 @@ public KafkaIndexTaskClient build( ) { return new KafkaIndexTaskClient( - httpClient, - mapper, + getHttpClient(), + getMapper(), taskInfoProvider, dataSource, numThreads, diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskIOConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskIOConfig.java new file mode 100644 index 000000000000..fc5c28751bc9 --- /dev/null +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskIOConfig.java @@ -0,0 +1,95 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import java.util.Map; + +public class KafkaIndexTaskIOConfig extends SeekableStreamIndexTaskIOConfig +{ + private final Map consumerProperties; + + @JsonCreator + public KafkaIndexTaskIOConfig( + @JsonProperty("taskGroupId") @Nullable Integer taskGroupId, // can be null for backward compabitility + @JsonProperty("baseSequenceName") String baseSequenceName, + @JsonProperty("startPartitions") SeekableStreamPartitions startPartitions, + @JsonProperty("endPartitions") SeekableStreamPartitions endPartitions, + @JsonProperty("consumerProperties") Map consumerProperties, + @JsonProperty("useTransaction") Boolean useTransaction, + @JsonProperty("minimumMessageTime") DateTime minimumMessageTime, + @JsonProperty("maximumMessageTime") DateTime maximumMessageTime, + @JsonProperty("skipOffsetGaps") Boolean skipOffsetGaps + ) + { + super( + taskGroupId, + baseSequenceName, + startPartitions, + endPartitions, + useTransaction, + minimumMessageTime, + maximumMessageTime, + skipOffsetGaps, + null + ); + + this.consumerProperties = Preconditions.checkNotNull(consumerProperties, "consumerProperties"); + + for (int partition : endPartitions.getPartitionSequenceNumberMap().keySet()) { + Preconditions.checkArgument( + endPartitions.getPartitionSequenceNumberMap() + .get(partition) + .compareTo(startPartitions.getPartitionSequenceNumberMap().get(partition)) >= 0, + "end offset must be >= start offset for partition[%s]", + partition + ); + } + } + + @JsonProperty + public Map getConsumerProperties() + { + return consumerProperties; + } + + @Override + public String toString() + { + return "KafkaIndexTaskIOConfig{" + + "taskGroupId=" + getTaskGroupId() + + ", baseSequenceName='" + getBaseSequenceName() + '\'' + + ", startPartitions=" + getStartPartitions() + + ", endPartitions=" + getEndPartitions() + + ", consumerProperties=" + consumerProperties + + ", useTransaction=" + isUseTransaction() + + ", minimumMessageTime=" + getMinimumMessageTime() + + ", maximumMessageTime=" + getMaximumMessageTime() + + ", skipOffsetGaps=" + isSkipOffsetGaps() + + '}'; + } +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskModule.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskModule.java index 751421057ab0..47b5df82aebe 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskModule.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskModule.java @@ -26,9 +26,9 @@ import com.google.inject.Binder; import com.google.inject.TypeLiteral; import org.apache.druid.guice.LazySingleton; -import org.apache.druid.indexing.common.task.IndexTaskClientFactory; import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpec; import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; import org.apache.druid.initialization.DruidModule; import java.util.List; @@ -43,7 +43,7 @@ public List getJacksonModules() .registerSubtypes( new NamedType(KafkaIndexTask.class, "index_kafka"), new NamedType(KafkaDataSourceMetadata.class, "kafka"), - new NamedType(KafkaIOConfig.class, "kafka"), + new NamedType(KafkaIndexTaskIOConfig.class, "kafka"), new NamedType(KafkaSupervisorTuningConfig.class, "kafka"), new NamedType(KafkaSupervisorSpec.class, "kafka") ) @@ -54,7 +54,7 @@ public List getJacksonModules() public void configure(Binder binder) { binder.bind( - new TypeLiteral>() + new TypeLiteral>() { } ).to(KafkaIndexTaskClientFactory.class).in(LazySingleton.class); diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskRunner.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskRunner.java deleted file mode 100644 index 348cfcc8c0b5..000000000000 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskRunner.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.druid.indexing.kafka; - -import com.google.common.annotations.VisibleForTesting; -import org.apache.druid.indexer.TaskStatus; -import org.apache.druid.indexing.common.TaskToolbox; -import org.apache.druid.indexing.common.stats.RowIngestionMeters; -import org.apache.druid.indexing.kafka.KafkaIndexTask.Status; -import org.apache.druid.segment.realtime.appenderator.Appenderator; -import org.apache.druid.segment.realtime.firehose.ChatHandler; - -import javax.ws.rs.core.Response; -import java.util.Map; - -/** - * This class is used by only {@link KafkaIndexTask}. We currently have two implementations of this interface, i.e., - * {@link IncrementalPublishingKafkaIndexTaskRunner} and {@link LegacyKafkaIndexTaskRunner}. The latter one was used in - * the versions prior to 0.12.0, but being kept to support rolling update from them. - * - * We don't have a good reason for having this interface except for better code maintenance for the latest kakfa - * indexing algorithm. As a result, this interface can be removed in the future when {@link LegacyKafkaIndexTaskRunner} - * is removed and it's no longer useful. - */ -public interface KafkaIndexTaskRunner extends ChatHandler -{ - Appenderator getAppenderator(); - - TaskStatus run(TaskToolbox toolbox); - - void stopGracefully(); - - // The below methods are mostly for unit testing. - - @VisibleForTesting - RowIngestionMeters getRowIngestionMeters(); - @VisibleForTesting - Status getStatus(); - - @VisibleForTesting - Map getCurrentOffsets(); - @VisibleForTesting - Map getEndOffsets(); - @VisibleForTesting - Response setEndOffsets( - Map offsets, - boolean finish // this field is only for internal purposes, shouldn't be usually set by users - ) throws InterruptedException; - - @VisibleForTesting - Response pause() throws InterruptedException; - @VisibleForTesting - void resume() throws InterruptedException; -} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfig.java new file mode 100644 index 000000000000..7cee87790654 --- /dev/null +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfig.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory; +import org.joda.time.Period; + +import javax.annotation.Nullable; +import java.io.File; + +@JsonTypeName("KafkaTuningConfig") +public class KafkaIndexTaskTuningConfig extends SeekableStreamIndexTaskTuningConfig +{ + @JsonCreator + public KafkaIndexTaskTuningConfig( + @JsonProperty("maxRowsInMemory") @Nullable Integer maxRowsInMemory, + @JsonProperty("maxBytesInMemory") @Nullable Long maxBytesInMemory, + @JsonProperty("maxRowsPerSegment") @Nullable Integer maxRowsPerSegment, + @JsonProperty("maxTotalRows") @Nullable Long maxTotalRows, + @JsonProperty("intermediatePersistPeriod") @Nullable Period intermediatePersistPeriod, + @JsonProperty("basePersistDirectory") @Nullable File basePersistDirectory, + @JsonProperty("maxPendingPersists") @Nullable Integer maxPendingPersists, + @JsonProperty("indexSpec") @Nullable IndexSpec indexSpec, + // This parameter is left for compatibility when reading existing configs, to be removed in Druid 0.12. + @JsonProperty("buildV9Directly") @Nullable Boolean buildV9Directly, + @Deprecated @JsonProperty("reportParseExceptions") @Nullable Boolean reportParseExceptions, + @JsonProperty("handoffConditionTimeout") @Nullable Long handoffConditionTimeout, + @JsonProperty("resetOffsetAutomatically") @Nullable Boolean resetOffsetAutomatically, + @JsonProperty("segmentWriteOutMediumFactory") @Nullable SegmentWriteOutMediumFactory segmentWriteOutMediumFactory, + @JsonProperty("intermediateHandoffPeriod") @Nullable Period intermediateHandoffPeriod, + @JsonProperty("logParseExceptions") @Nullable Boolean logParseExceptions, + @JsonProperty("maxParseExceptions") @Nullable Integer maxParseExceptions, + @JsonProperty("maxSavedParseExceptions") @Nullable Integer maxSavedParseExceptions + ) + { + super( + maxRowsInMemory, + maxBytesInMemory, + maxRowsPerSegment, + maxTotalRows, + intermediatePersistPeriod, + basePersistDirectory, + maxPendingPersists, + indexSpec, + true, + reportParseExceptions, + handoffConditionTimeout, + resetOffsetAutomatically, + false, + segmentWriteOutMediumFactory, + intermediateHandoffPeriod, + logParseExceptions, + maxParseExceptions, + maxSavedParseExceptions + ); + } + + @Override + public KafkaIndexTaskTuningConfig withBasePersistDirectory(File dir) + { + return new KafkaIndexTaskTuningConfig( + getMaxRowsInMemory(), + getMaxBytesInMemory(), + getMaxRowsPerSegment(), + getMaxTotalRows(), + getIntermediatePersistPeriod(), + dir, + getMaxPendingPersists(), + getIndexSpec(), + true, + isReportParseExceptions(), + getHandoffConditionTimeout(), + isResetOffsetAutomatically(), + getSegmentWriteOutMediumFactory(), + getIntermediateHandoffPeriod(), + isLogParseExceptions(), + getMaxParseExceptions(), + getMaxSavedParseExceptions() + ); + } + + + @Override + public String toString() + { + return "KafkaIndexTaskTuningConfig{" + + "maxRowsInMemory=" + getMaxRowsInMemory() + + ", maxRowsPerSegment=" + getMaxRowsPerSegment() + + ", maxTotalRows=" + getMaxTotalRows() + + ", maxBytesInMemory=" + getMaxBytesInMemory() + + ", intermediatePersistPeriod=" + getIntermediatePersistPeriod() + + ", basePersistDirectory=" + getBasePersistDirectory() + + ", maxPendingPersists=" + getMaxPendingPersists() + + ", indexSpec=" + getIndexSpec() + + ", reportParseExceptions=" + isReportParseExceptions() + + ", handoffConditionTimeout=" + getHandoffConditionTimeout() + + ", resetOffsetAutomatically=" + isResetOffsetAutomatically() + + ", segmentWriteOutMediumFactory=" + getSegmentWriteOutMediumFactory() + + ", intermediateHandoffPeriod=" + getIntermediateHandoffPeriod() + + ", logParseExceptions=" + isLogParseExceptions() + + ", maxParseExceptions=" + getMaxParseExceptions() + + ", maxSavedParseExceptions=" + getMaxSavedParseExceptions() + + '}'; + } + +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaPartitions.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaPartitions.java deleted file mode 100644 index a5b36e9e8e7e..000000000000 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaPartitions.java +++ /dev/null @@ -1,98 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.druid.indexing.kafka; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableMap; -import org.apache.druid.java.util.common.StringUtils; - -import java.util.Map; -import java.util.Objects; - -public class KafkaPartitions -{ - private final String topic; - private final Map partitionOffsetMap; - - @JsonCreator - public KafkaPartitions( - @JsonProperty("topic") final String topic, - @JsonProperty("partitionOffsetMap") final Map partitionOffsetMap - ) - { - this.topic = topic; - this.partitionOffsetMap = ImmutableMap.copyOf(partitionOffsetMap); - - // Validate partitionOffsetMap - for (Map.Entry entry : partitionOffsetMap.entrySet()) { - Preconditions.checkArgument( - entry.getValue() >= 0, - StringUtils.format( - "partition[%d] offset[%d] invalid", - entry.getKey(), - entry.getValue() - ) - ); - } - } - - @JsonProperty - public String getTopic() - { - return topic; - } - - @JsonProperty - public Map getPartitionOffsetMap() - { - return partitionOffsetMap; - } - - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - KafkaPartitions that = (KafkaPartitions) o; - return Objects.equals(topic, that.topic) && - Objects.equals(partitionOffsetMap, that.partitionOffsetMap); - } - - @Override - public int hashCode() - { - return Objects.hash(topic, partitionOffsetMap); - } - - @Override - public String toString() - { - return "KafkaPartitions{" + - "topic='" + topic + '\'' + - ", partitionOffsetMap=" + partitionOffsetMap + - '}'; - } -} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java new file mode 100644 index 000000000000..935404cbc7c2 --- /dev/null +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java @@ -0,0 +1,214 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorIOConfig; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.indexing.seekablestream.utils.RandomIdUtils; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.metadata.PasswordProvider; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.common.PartitionInfo; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; + +import javax.annotation.Nonnull; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.stream.Collectors; + +public class KafkaRecordSupplier implements RecordSupplier +{ + private final KafkaConsumer consumer; + private final Map consumerProperties; + private final ObjectMapper sortingMapper; + private boolean closed; + + public KafkaRecordSupplier( + Map consumerProperties, + ObjectMapper sortingMapper + ) + { + this.consumerProperties = consumerProperties; + this.sortingMapper = sortingMapper; + this.consumer = getKafkaConsumer(); + } + + @Override + public void assign(Set> streamPartitions) + { + consumer.assign(streamPartitions + .stream() + .map(x -> new TopicPartition(x.getStream(), x.getPartitionId())) + .collect(Collectors.toSet())); + seekToEarliest(streamPartitions); + } + + @Override + public void seek(StreamPartition partition, Long sequenceNumber) + { + consumer.seek(new TopicPartition(partition.getStream(), partition.getPartitionId()), sequenceNumber); + } + + @Override + public void seekToEarliest(Set> partitions) + { + consumer.seekToBeginning(partitions + .stream() + .map(e -> new TopicPartition(e.getStream(), e.getPartitionId())) + .collect(Collectors.toList())); + } + + @Override + public void seekToLatest(Set> partitions) + { + consumer.seekToEnd(partitions + .stream() + .map(e -> new TopicPartition(e.getStream(), e.getPartitionId())) + .collect(Collectors.toList())); + } + + @Override + public Set> getAssignment() + { + Set topicPartitions = consumer.assignment(); + return topicPartitions + .stream() + .map(e -> new StreamPartition<>(e.topic(), e.partition())) + .collect(Collectors.toSet()); + } + + @Nonnull + @Override + public List> poll(long timeout) + { + List> polledRecords = new ArrayList<>(); + for (ConsumerRecord record : consumer.poll(timeout)) { + polledRecords.add(new OrderedPartitionableRecord<>( + record.topic(), + record.partition(), + record.offset(), + record.value() == null ? null : ImmutableList.of(record.value()) + )); + } + return polledRecords; + } + + @Override + public Long getLatestSequenceNumber(StreamPartition partition) + { + Long currPos = consumer.position(new TopicPartition(partition.getStream(), partition.getPartitionId())); + seekToLatest(Collections.singleton(partition)); + Long nextPos = consumer.position(new TopicPartition(partition.getStream(), partition.getPartitionId())); + seek(partition, currPos); + return nextPos; + } + + @Override + public Long getEarliestSequenceNumber(StreamPartition partition) + { + Long currPos = consumer.position(new TopicPartition(partition.getStream(), partition.getPartitionId())); + seekToEarliest(Collections.singleton(partition)); + Long nextPos = consumer.position(new TopicPartition(partition.getStream(), partition.getPartitionId())); + seek(partition, currPos); + return nextPos; + } + + @Override + public Long getPosition(StreamPartition partition) + { + return consumer.position(new TopicPartition(partition.getStream(), partition.getPartitionId())); + } + + @Override + public Set getPartitionIds(String stream) + { + List partitions = consumer.partitionsFor(stream); + if (partitions == null) { + throw new ISE("Topic [%s] is not found in KafkaConsumer's list of topics", stream); + } + return partitions.stream().map(PartitionInfo::partition).collect(Collectors.toSet()); + } + + @Override + public void close() + { + if (closed) { + return; + } + closed = true; + consumer.close(); + } + + public static void addConsumerPropertiesFromConfig( + Properties properties, + ObjectMapper configMapper, + Map consumerProperties + ) + { + // Extract passwords before SSL connection to Kafka + for (Map.Entry entry : consumerProperties.entrySet()) { + String propertyKey = entry.getKey(); + if (propertyKey.equals(KafkaSupervisorIOConfig.TRUST_STORE_PASSWORD_KEY) + || propertyKey.equals(KafkaSupervisorIOConfig.KEY_STORE_PASSWORD_KEY) + || propertyKey.equals(KafkaSupervisorIOConfig.KEY_PASSWORD_KEY)) { + PasswordProvider configPasswordProvider = configMapper.convertValue( + entry.getValue(), + PasswordProvider.class + ); + properties.setProperty(propertyKey, configPasswordProvider.getPassword()); + } else { + properties.setProperty(propertyKey, String.valueOf(entry.getValue())); + } + } + } + + private KafkaConsumer getKafkaConsumer() + { + final Properties props = new Properties(); + + props.setProperty("metadata.max.age.ms", "10000"); + props.setProperty("group.id", StringUtils.format("kafka-supervisor-%s", RandomIdUtils.getRandomId())); + + addConsumerPropertiesFromConfig(props, sortingMapper, consumerProperties); + + props.setProperty("enable.auto.commit", "false"); + + ClassLoader currCtxCl = Thread.currentThread().getContextClassLoader(); + try { + Thread.currentThread().setContextClassLoader(getClass().getClassLoader()); + return new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer()); + } + finally { + Thread.currentThread().setContextClassLoader(currCtxCl); + } + } + +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSequenceNumber.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSequenceNumber.java new file mode 100644 index 000000000000..e9033521b07e --- /dev/null +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSequenceNumber.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; + +import javax.validation.constraints.NotNull; + +public class KafkaSequenceNumber extends OrderedSequenceNumber +{ + private KafkaSequenceNumber(Long sequenceNumber) + { + super(sequenceNumber, false); + } + + public static KafkaSequenceNumber of(Long sequenceNumber) + { + return new KafkaSequenceNumber(sequenceNumber); + } + + @Override + public int compareTo( + @NotNull OrderedSequenceNumber o + ) + { + return this.get().compareTo(o.get()); + } + +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/LegacyKafkaIndexTaskRunner.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/LegacyKafkaIndexTaskRunner.java index 75d268842497..53320c664da5 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/LegacyKafkaIndexTaskRunner.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/LegacyKafkaIndexTaskRunner.java @@ -45,7 +45,14 @@ import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; import org.apache.druid.indexing.common.task.IndexTaskUtils; import org.apache.druid.indexing.common.task.RealtimeIndexTask; -import org.apache.druid.indexing.kafka.KafkaIndexTask.Status; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; @@ -74,6 +81,7 @@ import org.apache.kafka.common.TopicPartition; import org.joda.time.DateTime; +import javax.annotation.Nonnull; import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; import javax.ws.rs.Consumes; @@ -93,6 +101,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.TreeMap; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; import java.util.concurrent.RejectedExecutionException; @@ -109,7 +118,7 @@ * This class will be removed in a future release. */ @Deprecated -public class LegacyKafkaIndexTaskRunner implements KafkaIndexTaskRunner +public class LegacyKafkaIndexTaskRunner extends SeekableStreamIndexTaskRunner { private static final EmittingLogger log = new EmittingLogger(LegacyKafkaIndexTaskRunner.class); private static final String METADATA_NEXT_PARTITIONS = "nextPartitions"; @@ -153,8 +162,8 @@ public class LegacyKafkaIndexTaskRunner implements KafkaIndexTaskRunner private final Condition isAwaitingRetry = pollRetryLock.newCondition(); private final KafkaIndexTask task; - private final KafkaIOConfig ioConfig; - private final KafkaTuningConfig tuningConfig; + private final KafkaIndexTaskIOConfig ioConfig; + private final KafkaIndexTaskTuningConfig tuningConfig; private final InputRowParser parser; private final AuthorizerMapper authorizerMapper; private final Optional chatHandlerProvider; @@ -181,6 +190,14 @@ public class LegacyKafkaIndexTaskRunner implements KafkaIndexTaskRunner RowIngestionMetersFactory rowIngestionMetersFactory ) { + super( + task, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory + ); this.task = task; this.ioConfig = task.getIOConfig(); this.tuningConfig = task.getTuningConfig(); @@ -190,7 +207,7 @@ public class LegacyKafkaIndexTaskRunner implements KafkaIndexTaskRunner this.savedParseExceptions = savedParseExceptions; this.rowIngestionMeters = rowIngestionMetersFactory.createRowIngestionMeters(); - this.endOffsets.putAll(ioConfig.getEndPartitions().getPartitionOffsetMap()); + this.endOffsets.putAll(ioConfig.getEndPartitions().getPartitionSequenceNumberMap()); this.ingestionState = IngestionState.NOT_STARTED; } @@ -274,34 +291,39 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception appenderator = appenderator0; - final String topic = ioConfig.getStartPartitions().getTopic(); + final String topic = ioConfig.getStartPartitions().getStream(); // Start up, set up initial offsets. final Object restoredMetadata = driver.startJob(); if (restoredMetadata == null) { - nextOffsets.putAll(ioConfig.getStartPartitions().getPartitionOffsetMap()); + nextOffsets.putAll(ioConfig.getStartPartitions().getPartitionSequenceNumberMap()); } else { final Map restoredMetadataMap = (Map) restoredMetadata; - final KafkaPartitions restoredNextPartitions = toolbox.getObjectMapper().convertValue( + final SeekableStreamPartitions restoredNextPartitions = toolbox.getObjectMapper().convertValue( restoredMetadataMap.get(METADATA_NEXT_PARTITIONS), - KafkaPartitions.class + toolbox.getObjectMapper().getTypeFactory().constructParametrizedType( + SeekableStreamPartitions.class, + SeekableStreamPartitions.class, + Integer.class, + Long.class + ) ); - nextOffsets.putAll(restoredNextPartitions.getPartitionOffsetMap()); + nextOffsets.putAll(restoredNextPartitions.getPartitionSequenceNumberMap()); // Sanity checks. - if (!restoredNextPartitions.getTopic().equals(ioConfig.getStartPartitions().getTopic())) { + if (!restoredNextPartitions.getStream().equals(ioConfig.getStartPartitions().getStream())) { throw new ISE( "WTF?! Restored topic[%s] but expected topic[%s]", - restoredNextPartitions.getTopic(), - ioConfig.getStartPartitions().getTopic() + restoredNextPartitions.getStream(), + ioConfig.getStartPartitions().getStream() ); } - if (!nextOffsets.keySet().equals(ioConfig.getStartPartitions().getPartitionOffsetMap().keySet())) { + if (!nextOffsets.keySet().equals(ioConfig.getStartPartitions().getPartitionSequenceNumberMap().keySet())) { throw new ISE( "WTF?! Restored partitions[%s] but expected partitions[%s]", nextOffsets.keySet(), - ioConfig.getStartPartitions().getPartitionOffsetMap().keySet() + ioConfig.getStartPartitions().getPartitionSequenceNumberMap().keySet() ); } } @@ -326,8 +348,8 @@ public Committer get() public Object getMetadata() { return ImmutableMap.of( - METADATA_NEXT_PARTITIONS, new KafkaPartitions( - ioConfig.getStartPartitions().getTopic(), + METADATA_NEXT_PARTITIONS, new SeekableStreamPartitions<>( + ioConfig.getStartPartitions().getStream(), snapshot ) ); @@ -494,13 +516,20 @@ public void run() } final TransactionalSegmentPublisher publisher = (segments, commitMetadata) -> { - final KafkaPartitions finalPartitions = toolbox.getObjectMapper().convertValue( + final SeekableStreamPartitions finalPartitions = toolbox.getObjectMapper().convertValue( ((Map) Preconditions.checkNotNull(commitMetadata, "commitMetadata")).get(METADATA_NEXT_PARTITIONS), - KafkaPartitions.class + toolbox.getObjectMapper() + .getTypeFactory() + .constructParametrizedType( + SeekableStreamPartitions.class, + SeekableStreamPartitions.class, + Integer.class, + Long.class + ) ); // Sanity check, we should only be publishing things that match our desired end state. - if (!endOffsets.equals(finalPartitions.getPartitionOffsetMap())) { + if (!endOffsets.equals(finalPartitions.getPartitionSequenceNumberMap())) { throw new ISE("WTF?! Driver attempted to publish invalid metadata[%s].", commitMetadata); } @@ -530,9 +559,9 @@ public void run() ).get(); final List publishedSegments = published.getSegments() - .stream() - .map(DataSegment::getIdentifier) - .collect(Collectors.toList()); + .stream() + .map(DataSegment::getIdentifier) + .collect(Collectors.toList()); log.info( "Published segments[%s] with metadata[%s].", @@ -596,6 +625,21 @@ public void run() ); } + @Override + protected boolean isEndOfShard(Long seqNum) + { + return false; + } + + @Nonnull + @Override + protected List> getRecords( + RecordSupplier recordSupplier, TaskToolbox toolbox + ) + { + throw new UnsupportedOperationException(); + } + private Set assignPartitionsAndSeekToNext(KafkaConsumer consumer, String topic) { // Initialize consumer assignment. @@ -615,7 +659,7 @@ private Set assignPartitionsAndSeekToNext(KafkaConsumer consumer, Strin } } - task.assignPartitions(consumer, topic, assignment); + KafkaIndexTask.assignPartitions(consumer, topic, assignment); // Seek to starting offsets. for (final int partition : assignment) { @@ -661,6 +705,39 @@ private boolean possiblyPause() throws InterruptedException return false; } + @Override + protected void possiblyResetDataSourceMetadata( + TaskToolbox toolbox, + RecordSupplier recordSupplier, + Set> assignment, + Map currOffsets + ) + { + throw new UnsupportedOperationException(); + } + + @Override + protected boolean isEndSequenceOffsetsExclusive() + { + return false; + } + + @Override + protected boolean isStartingSequenceOffsetsExclusive() + { + return false; + } + + + @Override + protected SeekableStreamPartitions deserializeSeekableStreamPartitionsFromMetadata( + ObjectMapper mapper, + Object object + ) + { + throw new UnsupportedOperationException(); + } + private void possiblyResetOffsetsOrWait( Map outOfRangePartitions, KafkaConsumer consumer, @@ -689,7 +766,7 @@ private void possiblyResetOffsetsOrWait( } if (doReset) { - sendResetRequestAndWait(resetPartitions, taskToolbox); + sendResetRequestAndWaitLegacy(resetPartitions, taskToolbox); } else { log.warn("Retrying in %dms", task.getPollRetryMs()); pollRetryLock.lockInterruptibly(); @@ -705,7 +782,7 @@ private void possiblyResetOffsetsOrWait( } } - private void sendResetRequestAndWait(Map outOfRangePartitions, TaskToolbox taskToolbox) + private void sendResetRequestAndWaitLegacy(Map outOfRangePartitions, TaskToolbox taskToolbox) throws IOException { Map partitionOffsetMap = new HashMap<>(); @@ -715,9 +792,9 @@ private void sendResetRequestAndWait(Map outOfRangePartiti boolean result = taskToolbox.getTaskActionClient() .submit(new ResetDataSourceMetadataAction( task.getDataSource(), - new KafkaDataSourceMetadata(new KafkaPartitions( + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( ioConfig.getStartPartitions() - .getTopic(), + .getStream(), partitionOffsetMap )) )); @@ -738,6 +815,12 @@ private void requestPause() pauseRequested = true; } + @Override + protected Long getSequenceNumberToStoreAfterRead(Long sequenceNumber) + { + throw new UnsupportedOperationException(); + } + private void handleParseException(ParseException pe, ConsumerRecord record) { if (pe.isFromPartiallyValidRow()) { @@ -817,7 +900,7 @@ public void stopGracefully() } try { - if (pauseLock.tryLock(KafkaIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + if (pauseLock.tryLock(SeekableStreamIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { try { if (pauseRequested) { pauseRequested = false; @@ -833,7 +916,7 @@ public void stopGracefully() return; } - if (pollRetryLock.tryLock(KafkaIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + if (pollRetryLock.tryLock(SeekableStreamIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { try { isAwaitingRetry.signalAll(); } @@ -860,6 +943,7 @@ private Access authorizationCheck(final HttpServletRequest req, Action action) return IndexTaskUtils.datasourceAuthorizationCheck(req, action, task.getDataSource(), authorizerMapper); } + @Override @POST @Path("/stop") public Response stop(@Context final HttpServletRequest req) @@ -869,6 +953,7 @@ public Response stop(@Context final HttpServletRequest req) return Response.status(Response.Status.OK).build(); } + @Override @GET @Path("/status") @Produces(MediaType.APPLICATION_JSON) @@ -884,6 +969,7 @@ public Status getStatus() return status; } + @Override @GET @Path("/offsets/current") @Produces(MediaType.APPLICATION_JSON) @@ -899,6 +985,7 @@ public Map getCurrentOffsets() return nextOffsets; } + @Override @GET @Path("/offsets/end") @Produces(MediaType.APPLICATION_JSON) @@ -915,10 +1002,10 @@ public Map getEndOffsets() } @Override - public Response setEndOffsets(Map offsets, boolean finish) throws InterruptedException + public Response setEndOffsets(Map sequenceNumbers, boolean finish) throws InterruptedException { // finish is not used in this mode - return setEndOffsets(offsets); + return setEndOffsets(sequenceNumbers); } @POST @@ -934,6 +1021,7 @@ public Response setEndOffsetsHTTP( return setEndOffsets(offsets); } + @Override @GET @Path("/rowStats") @Produces(MediaType.APPLICATION_JSON) @@ -960,6 +1048,7 @@ public Response getRowStats( return Response.ok(returnMap).build(); } + @Override @GET @Path("/unparseableEvents") @Produces(MediaType.APPLICATION_JSON) @@ -1037,6 +1126,7 @@ private boolean isPaused() * method has timed out and returned before the task has paused; 200 OK with a map of the current partition offsets * in the response body if the task successfully paused */ + @Override @POST @Path("/pause") @Produces(MediaType.APPLICATION_JSON) @@ -1095,6 +1185,7 @@ public Response pause() throws InterruptedException } } + @Override @POST @Path("/resume") public Response resumeHTTP(@Context final HttpServletRequest req) throws InterruptedException @@ -1125,6 +1216,21 @@ public void resume() throws InterruptedException } } + @Override + protected SeekableStreamDataSourceMetadata createDataSourceMetadata( + SeekableStreamPartitions partitions + ) + { + throw new UnsupportedOperationException(); + } + + @Override + protected OrderedSequenceNumber createSequenceNumber(Long sequenceNumber) + { + throw new UnsupportedOperationException(); + } + + @Override @GET @Path("/time/start") @Produces(MediaType.APPLICATION_JSON) @@ -1133,4 +1239,15 @@ public DateTime getStartTime(@Context final HttpServletRequest req) authorizationCheck(req, Action.WRITE); return startTime; } + + @Nullable + @Override + protected TreeMap> getCheckPointsFromContext( + TaskToolbox toolbox, + String checkpointsString + ) + { + throw new UnsupportedOperationException(); + } + } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java index b0c4ba56ec55..6f4434ea3e95 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java @@ -21,95 +21,52 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.MapperFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Function; import com.google.common.base.Joiner; -import com.google.common.base.Optional; -import com.google.common.base.Preconditions; -import com.google.common.base.Predicate; -import com.google.common.base.Throwables; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Iterables; -import com.google.common.primitives.Longs; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.ListeningExecutorService; -import com.google.common.util.concurrent.MoreExecutors; -import org.apache.commons.codec.digest.DigestUtils; -import org.apache.druid.indexer.TaskLocation; -import org.apache.druid.indexer.TaskStatus; -import org.apache.druid.indexing.common.IndexTaskClient; -import org.apache.druid.indexing.common.TaskInfoProvider; import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; -import org.apache.druid.indexing.common.task.RealtimeIndexTask; import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.common.task.TaskResource; import org.apache.druid.indexing.kafka.KafkaDataSourceMetadata; -import org.apache.druid.indexing.kafka.KafkaIOConfig; import org.apache.druid.indexing.kafka.KafkaIndexTask; -import org.apache.druid.indexing.kafka.KafkaIndexTaskClient; import org.apache.druid.indexing.kafka.KafkaIndexTaskClientFactory; -import org.apache.druid.indexing.kafka.KafkaPartitions; -import org.apache.druid.indexing.kafka.KafkaTuningConfig; +import org.apache.druid.indexing.kafka.KafkaIndexTaskIOConfig; +import org.apache.druid.indexing.kafka.KafkaIndexTaskTuningConfig; +import org.apache.druid.indexing.kafka.KafkaRecordSupplier; +import org.apache.druid.indexing.kafka.KafkaSequenceNumber; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; import org.apache.druid.indexing.overlord.TaskMaster; -import org.apache.druid.indexing.overlord.TaskQueue; -import org.apache.druid.indexing.overlord.TaskRunner; -import org.apache.druid.indexing.overlord.TaskRunnerListener; -import org.apache.druid.indexing.overlord.TaskRunnerWorkItem; import org.apache.druid.indexing.overlord.TaskStorage; -import org.apache.druid.indexing.overlord.supervisor.Supervisor; -import org.apache.druid.indexing.overlord.supervisor.SupervisorReport; -import org.apache.druid.java.util.common.DateTimes; -import org.apache.druid.java.util.common.IAE; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorReportPayload; +import org.apache.druid.indexing.seekablestream.utils.RandomIdUtils; import org.apache.druid.java.util.common.ISE; -import org.apache.druid.java.util.common.Pair; -import org.apache.druid.java.util.common.RetryUtils; import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.emitter.EmittingLogger; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.java.util.emitter.service.ServiceMetricEvent; -import org.apache.druid.metadata.EntryExistsException; import org.apache.druid.server.metrics.DruidMonitorSchedulerConfig; -import org.apache.kafka.clients.consumer.KafkaConsumer; -import org.apache.kafka.common.PartitionInfo; -import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.joda.time.DateTime; -import javax.annotation.Nullable; -import java.io.IOException; import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Map.Entry; -import java.util.Properties; import java.util.Set; -import java.util.SortedMap; import java.util.TreeMap; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.LinkedBlockingDeque; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; -import java.util.stream.Stream; /** * Supervisor responsible for managing the KafkaIndexTasks for a single dataSource. At a high level, the class accepts a @@ -119,2228 +76,194 @@ * tasks to satisfy the desired number of replicas. As tasks complete, new tasks are queued to process the next range of * Kafka offsets. */ -public class KafkaSupervisor implements Supervisor +public class KafkaSupervisor extends SeekableStreamSupervisor { private static final EmittingLogger log = new EmittingLogger(KafkaSupervisor.class); - private static final long MAX_RUN_FREQUENCY_MILLIS = 1000; // prevent us from running too often in response to events - private static final long NOT_SET = -1; - private static final long MINIMUM_FUTURE_TIMEOUT_IN_SECONDS = 120; private static final long MINIMUM_GET_OFFSET_PERIOD_MILLIS = 5000; private static final long INITIAL_GET_OFFSET_DELAY_MILLIS = 15000; private static final long INITIAL_EMIT_LAG_METRIC_DELAY_MILLIS = 25000; - private static final int MAX_INITIALIZATION_RETRIES = 20; + private static final Long NOT_SET = -1L; + private static final Long END_OF_PARTITION = Long.MAX_VALUE; - public static final String IS_INCREMENTAL_HANDOFF_SUPPORTED = "IS_INCREMENTAL_HANDOFF_SUPPORTED"; - - // Internal data structures - // -------------------------------------------------------- - - /** - * A TaskGroup is the main data structure used by KafkaSupervisor to organize and monitor Kafka partitions and - * indexing tasks. All the tasks in a TaskGroup should always be doing the same thing (reading the same partitions and - * starting from the same offset) and if [replicas] is configured to be 1, a TaskGroup will contain a single task (the - * exception being if the supervisor started up and discovered and adopted some already running tasks). At any given - * time, there should only be up to a maximum of [taskCount] actively-reading task groups (tracked in the [taskGroups] - * map) + zero or more pending-completion task groups (tracked in [pendingCompletionTaskGroups]). - */ - private class TaskGroup - { - final int groupId; - - // This specifies the partitions and starting offsets for this task group. It is set on group creation from the data - // in [partitionGroups] and never changes during the lifetime of this task group, which will live until a task in - // this task group has completed successfully, at which point this will be destroyed and a new task group will be - // created with new starting offsets. This allows us to create replacement tasks for failed tasks that process the - // same offsets, even if the values in [partitionGroups] has been changed. - final ImmutableMap partitionOffsets; - - final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); - final Optional minimumMessageTime; - final Optional maximumMessageTime; - DateTime completionTimeout; // is set after signalTasksToFinish(); if not done by timeout, take corrective action - final TreeMap> sequenceOffsets = new TreeMap<>(); - final String baseSequenceName; - - TaskGroup( - int groupId, - ImmutableMap partitionOffsets, - Optional minimumMessageTime, - Optional maximumMessageTime - ) - { - this.groupId = groupId; - this.partitionOffsets = partitionOffsets; - this.minimumMessageTime = minimumMessageTime; - this.maximumMessageTime = maximumMessageTime; - this.sequenceOffsets.put(0, partitionOffsets); - this.baseSequenceName = generateSequenceName(partitionOffsets, minimumMessageTime, maximumMessageTime); - } - - int addNewCheckpoint(Map checkpoint) - { - sequenceOffsets.put(sequenceOffsets.lastKey() + 1, checkpoint); - return sequenceOffsets.lastKey(); - } - - Set taskIds() - { - return tasks.keySet(); - } - } - - private static class TaskData - { - @Nullable - volatile TaskStatus status; - @Nullable - volatile DateTime startTime; - volatile Map currentOffsets = new HashMap<>(); - - @Override - public String toString() - { - return "TaskData{" + - "status=" + status + - ", startTime=" + startTime + - ", currentOffsets=" + currentOffsets + - '}'; - } - } - - // Map<{group ID}, {actively reading task group}>; see documentation for TaskGroup class - private final ConcurrentHashMap taskGroups = new ConcurrentHashMap<>(); - - // After telling a taskGroup to stop reading and begin publishing a segment, it is moved from [taskGroups] to here so - // we can monitor its status while we queue new tasks to read the next range of offsets. This is a list since we could - // have multiple sets of tasks publishing at once if time-to-publish > taskDuration. - // Map<{group ID}, List<{pending completion task groups}>> - private final ConcurrentHashMap> pendingCompletionTaskGroups = new ConcurrentHashMap<>(); - - // The starting offset for a new partition in [partitionGroups] is initially set to NOT_SET. When a new task group - // is created and is assigned partitions, if the offset in [partitionGroups] is NOT_SET it will take the starting - // offset value from the metadata store, and if it can't find it there, from Kafka. Once a task begins - // publishing, the offset in partitionGroups will be updated to the ending offset of the publishing-but-not-yet- - // completed task, which will cause the next set of tasks to begin reading from where the previous task left - // off. If that previous task now fails, we will set the offset in [partitionGroups] back to NOT_SET which will - // cause successive tasks to again grab their starting offset from metadata store. This mechanism allows us to - // start up successive tasks without waiting for the previous tasks to succeed and still be able to handle task - // failures during publishing. - // Map<{group ID}, Map<{partition ID}, {startingOffset}>> - private final ConcurrentHashMap> partitionGroups = new ConcurrentHashMap<>(); - // -------------------------------------------------------- - - private final TaskStorage taskStorage; - private final TaskMaster taskMaster; - private final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator; - private final KafkaIndexTaskClient taskClient; - private final ObjectMapper sortingMapper; - private final KafkaSupervisorSpec spec; - private final ServiceEmitter emitter; - private final DruidMonitorSchedulerConfig monitorSchedulerConfig; - private final String dataSource; - private final KafkaSupervisorIOConfig ioConfig; - private final KafkaSupervisorTuningConfig tuningConfig; - private final KafkaTuningConfig taskTuningConfig; - private final String supervisorId; - private final TaskInfoProvider taskInfoProvider; - private final long futureTimeoutInSeconds; // how long to wait for async operations to complete - private final RowIngestionMetersFactory rowIngestionMetersFactory; - - private final ExecutorService exec; - private final ScheduledExecutorService scheduledExec; - private final ScheduledExecutorService reportingExec; - private final ListeningExecutorService workerExec; - private final BlockingQueue notices = new LinkedBlockingDeque<>(); - private final Object stopLock = new Object(); - private final Object stateChangeLock = new Object(); - private final Object consumerLock = new Object(); - - private boolean listenerRegistered = false; - private long lastRunTime; - - private int initRetryCounter = 0; - - private volatile DateTime firstRunTime; - private volatile KafkaConsumer consumer; - - private volatile boolean lifecycleStarted = false; - private volatile boolean started = false; - private volatile boolean stopped = false; - private volatile Map latestOffsetsFromKafka; - private volatile DateTime offsetsLastUpdated; - - public KafkaSupervisor( - final TaskStorage taskStorage, - final TaskMaster taskMaster, - final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, - final KafkaIndexTaskClientFactory taskClientFactory, - final ObjectMapper mapper, - final KafkaSupervisorSpec spec, - final RowIngestionMetersFactory rowIngestionMetersFactory - ) - { - this.taskStorage = taskStorage; - this.taskMaster = taskMaster; - this.indexerMetadataStorageCoordinator = indexerMetadataStorageCoordinator; - this.sortingMapper = mapper.copy().configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true); - this.spec = spec; - this.emitter = spec.getEmitter(); - this.monitorSchedulerConfig = spec.getMonitorSchedulerConfig(); - this.rowIngestionMetersFactory = rowIngestionMetersFactory; - - this.dataSource = spec.getDataSchema().getDataSource(); - this.ioConfig = spec.getIoConfig(); - this.tuningConfig = spec.getTuningConfig(); - this.taskTuningConfig = KafkaTuningConfig.copyOf(this.tuningConfig); - this.supervisorId = StringUtils.format("KafkaSupervisor-%s", dataSource); - this.exec = Execs.singleThreaded(supervisorId); - this.scheduledExec = Execs.scheduledSingleThreaded(supervisorId + "-Scheduler-%d"); - this.reportingExec = Execs.scheduledSingleThreaded(supervisorId + "-Reporting-%d"); - - int workerThreads = (this.tuningConfig.getWorkerThreads() != null - ? this.tuningConfig.getWorkerThreads() - : Math.min(10, this.ioConfig.getTaskCount())); - this.workerExec = MoreExecutors.listeningDecorator(Execs.multiThreaded(workerThreads, supervisorId + "-Worker-%d")); - log.info("Created worker pool with [%d] threads for dataSource [%s]", workerThreads, this.dataSource); - - this.taskInfoProvider = new TaskInfoProvider() - { - @Override - public TaskLocation getTaskLocation(final String id) - { - Preconditions.checkNotNull(id, "id"); - Optional taskRunner = taskMaster.getTaskRunner(); - if (taskRunner.isPresent()) { - Optional item = Iterables.tryFind( - taskRunner.get().getRunningTasks(), - (Predicate) taskRunnerWorkItem -> id.equals(taskRunnerWorkItem.getTaskId()) - ); - - if (item.isPresent()) { - return item.get().getLocation(); - } - } else { - log.error("Failed to get task runner because I'm not the leader!"); - } - - return TaskLocation.unknown(); - } - - @Override - public Optional getTaskStatus(String id) - { - return taskStorage.getStatus(id); - } - }; - - this.futureTimeoutInSeconds = Math.max( - MINIMUM_FUTURE_TIMEOUT_IN_SECONDS, - tuningConfig.getChatRetries() * (tuningConfig.getHttpTimeout().getStandardSeconds() - + IndexTaskClient.MAX_RETRY_WAIT_SECONDS) - ); - - int chatThreads = (this.tuningConfig.getChatThreads() != null - ? this.tuningConfig.getChatThreads() - : Math.min(10, this.ioConfig.getTaskCount() * this.ioConfig.getReplicas())); - this.taskClient = taskClientFactory.build( - taskInfoProvider, - dataSource, - chatThreads, - this.tuningConfig.getHttpTimeout(), - this.tuningConfig.getChatRetries() - ); - log.info( - "Created taskClient with dataSource[%s] chatThreads[%d] httpTimeout[%s] chatRetries[%d]", - dataSource, - chatThreads, - this.tuningConfig.getHttpTimeout(), - this.tuningConfig.getChatRetries() - ); - } - - @Override - public void start() - { - synchronized (stateChangeLock) { - Preconditions.checkState(!lifecycleStarted, "already started"); - Preconditions.checkState(!exec.isShutdown(), "already stopped"); - - // Try normal initialization first, if that fails then schedule periodic initialization retries - try { - tryInit(); - } - catch (Exception e) { - if (!started) { - log.warn("First initialization attempt failed for KafkaSupervisor[%s], starting retries...", dataSource); - - exec.submit( - () -> { - try { - RetryUtils.retry( - () -> { - tryInit(); - return 0; - }, - (throwable) -> { - return !started; - }, - 0, - MAX_INITIALIZATION_RETRIES, - null, - null - ); - } - catch (Exception e2) { - log.makeAlert( - "Failed to initialize after %s retries, aborting. Please resubmit the supervisor spec to restart this supervisor [%s]", - MAX_INITIALIZATION_RETRIES, - supervisorId - ).emit(); - throw new RuntimeException(e2); - } - } - ); - } - } - - lifecycleStarted = true; - } - } - - @Override - public void stop(boolean stopGracefully) - { - synchronized (stateChangeLock) { - Preconditions.checkState(lifecycleStarted, "lifecycle not started"); - - log.info("Beginning shutdown of KafkaSupervisor[%s]", dataSource); - - try { - scheduledExec.shutdownNow(); // stop recurring executions - reportingExec.shutdownNow(); - - if (started) { - Optional taskRunner = taskMaster.getTaskRunner(); - if (taskRunner.isPresent()) { - taskRunner.get().unregisterListener(supervisorId); - } - - // Stopping gracefully will synchronize the end offsets of the tasks and signal them to publish, and will block - // until the tasks have acknowledged or timed out. We want this behavior when we're explicitly shut down through - // the API, but if we shut down for other reasons (e.g. we lose leadership) we want to just stop and leave the - // tasks as they are. - synchronized (stopLock) { - if (stopGracefully) { - log.info("Posting GracefulShutdownNotice, signalling managed tasks to complete and publish"); - notices.add(new GracefulShutdownNotice()); - } else { - log.info("Posting ShutdownNotice"); - notices.add(new ShutdownNotice()); - } - - long shutdownTimeoutMillis = tuningConfig.getShutdownTimeout().getMillis(); - long endTime = System.currentTimeMillis() + shutdownTimeoutMillis; - while (!stopped) { - long sleepTime = endTime - System.currentTimeMillis(); - if (sleepTime <= 0) { - log.info("Timed out while waiting for shutdown (timeout [%,dms])", shutdownTimeoutMillis); - stopped = true; - break; - } - stopLock.wait(sleepTime); - } - } - log.info("Shutdown notice handled"); - } - - taskClient.close(); - workerExec.shutdownNow(); - exec.shutdownNow(); - started = false; - - log.info("KafkaSupervisor[%s] has stopped", dataSource); - } - catch (Exception e) { - log.makeAlert(e, "Exception stopping KafkaSupervisor[%s]", dataSource) - .emit(); - } - } - } - - private boolean someTaskGroupsPendingCompletion(Integer groupId) - { - CopyOnWriteArrayList taskGroups = pendingCompletionTaskGroups.get(groupId); - return taskGroups != null && taskGroups.size() > 0; - } - - @Override - public SupervisorReport getStatus() - { - return generateReport(true); - } - - @Override - public Map> getStats() - { - try { - return getCurrentTotalStats(); - } - catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - log.error(ie, "getStats() interrupted."); - throw new RuntimeException(ie); - } - catch (ExecutionException | TimeoutException eete) { - throw new RuntimeException(eete); - } - } - - @Override - public void reset(DataSourceMetadata dataSourceMetadata) - { - log.info("Posting ResetNotice"); - notices.add(new ResetNotice(dataSourceMetadata)); - } - - @Override - public void checkpoint( - @Nullable Integer taskGroupId, - @Deprecated String baseSequenceName, - DataSourceMetadata previousCheckPoint, - DataSourceMetadata currentCheckPoint - ) - { - Preconditions.checkNotNull(previousCheckPoint, "previousCheckpoint"); - Preconditions.checkNotNull(currentCheckPoint, "current checkpoint cannot be null"); - Preconditions.checkArgument( - ioConfig.getTopic().equals(((KafkaDataSourceMetadata) currentCheckPoint).getKafkaPartitions().getTopic()), - "Supervisor topic [%s] and topic in checkpoint [%s] does not match", - ioConfig.getTopic(), - ((KafkaDataSourceMetadata) currentCheckPoint).getKafkaPartitions().getTopic() - ); - - log.info("Checkpointing [%s] for taskGroup [%s]", currentCheckPoint, taskGroupId); - notices.add( - new CheckpointNotice( - taskGroupId, - baseSequenceName, - (KafkaDataSourceMetadata) previousCheckPoint, - (KafkaDataSourceMetadata) currentCheckPoint - ) - ); - } - - public void possiblyRegisterListener() - { - // getTaskRunner() sometimes fails if the task queue is still being initialized so retry later until we succeed - - if (listenerRegistered) { - return; - } - - Optional taskRunner = taskMaster.getTaskRunner(); - if (taskRunner.isPresent()) { - taskRunner.get().registerListener( - new TaskRunnerListener() - { - @Override - public String getListenerId() - { - return supervisorId; - } - - @Override - public void locationChanged(final String taskId, final TaskLocation newLocation) - { - // do nothing - } - - @Override - public void statusChanged(String taskId, TaskStatus status) - { - notices.add(new RunNotice()); - } - }, MoreExecutors.sameThreadExecutor() - ); - - listenerRegistered = true; - } - } - - private interface Notice - { - void handle() throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException; - } - - private class RunNotice implements Notice - { - @Override - public void handle() throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException - { - long nowTime = System.currentTimeMillis(); - if (nowTime - lastRunTime < MAX_RUN_FREQUENCY_MILLIS) { - return; - } - lastRunTime = nowTime; - - runInternal(); - } - } - - private class GracefulShutdownNotice extends ShutdownNotice - { - @Override - public void handle() throws InterruptedException, ExecutionException, TimeoutException - { - gracefulShutdownInternal(); - super.handle(); - } - } - - private class ShutdownNotice implements Notice - { - @Override - public void handle() throws InterruptedException, ExecutionException, TimeoutException - { - consumer.close(); - - synchronized (stopLock) { - stopped = true; - stopLock.notifyAll(); - } - } - } - - private class ResetNotice implements Notice - { - final DataSourceMetadata dataSourceMetadata; - - ResetNotice(DataSourceMetadata dataSourceMetadata) - { - this.dataSourceMetadata = dataSourceMetadata; - } - - @Override - public void handle() - { - resetInternal(dataSourceMetadata); - } - } - - private class CheckpointNotice implements Notice - { - @Nullable - private final Integer nullableTaskGroupId; - @Deprecated - private final String baseSequenceName; - private final KafkaDataSourceMetadata previousCheckpoint; - private final KafkaDataSourceMetadata currentCheckpoint; - - CheckpointNotice( - @Nullable Integer nullableTaskGroupId, - @Deprecated String baseSequenceName, - KafkaDataSourceMetadata previousCheckpoint, - KafkaDataSourceMetadata currentCheckpoint - ) - { - this.baseSequenceName = baseSequenceName; - this.nullableTaskGroupId = nullableTaskGroupId; - this.previousCheckpoint = previousCheckpoint; - this.currentCheckpoint = currentCheckpoint; - } - - @Override - public void handle() throws ExecutionException, InterruptedException - { - // Find taskGroupId using taskId if it's null. It can be null while rolling update. - final int taskGroupId; - if (nullableTaskGroupId == null) { - // We search taskId in taskGroups and pendingCompletionTaskGroups sequentially. This should be fine because - // 1) a taskGroup can be moved from taskGroups to pendingCompletionTaskGroups in RunNotice - // (see checkTaskDuration()). - // 2) Notices are proceesed by a single thread. So, CheckpointNotice and RunNotice cannot be processed at the - // same time. - final java.util.Optional maybeGroupId = taskGroups - .entrySet() - .stream() - .filter(entry -> { - final TaskGroup taskGroup = entry.getValue(); - return taskGroup.baseSequenceName.equals(baseSequenceName); - }) - .findAny() - .map(Entry::getKey); - - if (maybeGroupId.isPresent()) { - taskGroupId = maybeGroupId.get(); - } else { - taskGroupId = pendingCompletionTaskGroups - .entrySet() - .stream() - .filter(entry -> { - final List taskGroups = entry.getValue(); - return taskGroups.stream().anyMatch(group -> group.baseSequenceName.equals(baseSequenceName)); - }) - .findAny() - .orElseThrow(() -> new ISE("Cannot find taskGroup for baseSequenceName[%s]", baseSequenceName)) - .getKey(); - } - } else { - taskGroupId = nullableTaskGroupId; - } - - // check for consistency - // if already received request for this sequenceName and dataSourceMetadata combination then return - final TaskGroup taskGroup = taskGroups.get(taskGroupId); - - if (isValidTaskGroup(taskGroupId, taskGroup)) { - final TreeMap> checkpoints = taskGroup.sequenceOffsets; - - // check validity of previousCheckpoint - int index = checkpoints.size(); - for (int sequenceId : checkpoints.descendingKeySet()) { - Map checkpoint = checkpoints.get(sequenceId); - // We have already verified the topic of the current checkpoint is same with that in ioConfig. - // See checkpoint(). - if (checkpoint.equals(previousCheckpoint.getKafkaPartitions().getPartitionOffsetMap())) { - break; - } - index--; - } - if (index == 0) { - throw new ISE("No such previous checkpoint [%s] found", previousCheckpoint); - } else if (index < checkpoints.size()) { - // if the found checkpoint is not the latest one then already checkpointed by a replica - Preconditions.checkState(index == checkpoints.size() - 1, "checkpoint consistency failure"); - log.info("Already checkpointed with offsets [%s]", checkpoints.lastEntry().getValue()); - return; - } - final Map newCheckpoint = checkpointTaskGroup(taskGroup, false).get(); - taskGroup.addNewCheckpoint(newCheckpoint); - log.info("Handled checkpoint notice, new checkpoint is [%s] for taskGroup [%s]", newCheckpoint, taskGroupId); - } - } - - private boolean isValidTaskGroup(int taskGroupId, @Nullable TaskGroup taskGroup) - { - if (taskGroup == null) { - // taskGroup might be in pendingCompletionTaskGroups or partitionGroups - if (pendingCompletionTaskGroups.containsKey(taskGroupId)) { - log.warn( - "Ignoring checkpoint request because taskGroup[%d] has already stopped indexing and is waiting for " - + "publishing segments", - taskGroupId - ); - return false; - } else if (partitionGroups.containsKey(taskGroupId)) { - log.warn("Ignoring checkpoint request because taskGroup[%d] is inactive", taskGroupId); - return false; - } else { - throw new ISE("WTH?! cannot find taskGroup [%s] among all taskGroups [%s]", taskGroupId, taskGroups); - } - } - - return true; - } - } - - @VisibleForTesting - void resetInternal(DataSourceMetadata dataSourceMetadata) - { - if (dataSourceMetadata == null) { - // Reset everything - boolean result = indexerMetadataStorageCoordinator.deleteDataSourceMetadata(dataSource); - log.info("Reset dataSource[%s] - dataSource metadata entry deleted? [%s]", dataSource, result); - taskGroups.values().forEach(group -> killTasksInGroup(group, "DataSourceMetadata is not found while reset")); - taskGroups.clear(); - partitionGroups.clear(); - } else if (!(dataSourceMetadata instanceof KafkaDataSourceMetadata)) { - throw new IAE("Expected KafkaDataSourceMetadata but found instance of [%s]", dataSourceMetadata.getClass()); - } else { - // Reset only the partitions in dataSourceMetadata if it has not been reset yet - final KafkaDataSourceMetadata resetKafkaMetadata = (KafkaDataSourceMetadata) dataSourceMetadata; - - if (resetKafkaMetadata.getKafkaPartitions().getTopic().equals(ioConfig.getTopic())) { - // metadata can be null - final DataSourceMetadata metadata = indexerMetadataStorageCoordinator.getDataSourceMetadata(dataSource); - if (metadata != null && !(metadata instanceof KafkaDataSourceMetadata)) { - throw new IAE( - "Expected KafkaDataSourceMetadata from metadata store but found instance of [%s]", - metadata.getClass() - ); - } - final KafkaDataSourceMetadata currentMetadata = (KafkaDataSourceMetadata) metadata; - - // defend against consecutive reset requests from replicas - // as well as the case where the metadata store do not have an entry for the reset partitions - boolean doReset = false; - for (Entry resetPartitionOffset : resetKafkaMetadata.getKafkaPartitions() - .getPartitionOffsetMap() - .entrySet()) { - final Long partitionOffsetInMetadataStore = currentMetadata == null - ? null - : currentMetadata.getKafkaPartitions() - .getPartitionOffsetMap() - .get(resetPartitionOffset.getKey()); - final TaskGroup partitionTaskGroup = taskGroups.get( - getTaskGroupIdForPartition(resetPartitionOffset.getKey()) - ); - final boolean isSameOffset = partitionTaskGroup != null - && partitionTaskGroup.partitionOffsets.get(resetPartitionOffset.getKey()) - .equals(resetPartitionOffset.getValue()); - if (partitionOffsetInMetadataStore != null || isSameOffset) { - doReset = true; - break; - } - } - - if (!doReset) { - log.info("Ignoring duplicate reset request [%s]", dataSourceMetadata); - return; - } - - boolean metadataUpdateSuccess = false; - if (currentMetadata == null) { - metadataUpdateSuccess = true; - } else { - final DataSourceMetadata newMetadata = currentMetadata.minus(resetKafkaMetadata); - try { - metadataUpdateSuccess = indexerMetadataStorageCoordinator.resetDataSourceMetadata(dataSource, newMetadata); - } - catch (IOException e) { - log.error("Resetting DataSourceMetadata failed [%s]", e.getMessage()); - Throwables.propagate(e); - } - } - if (metadataUpdateSuccess) { - resetKafkaMetadata.getKafkaPartitions().getPartitionOffsetMap().keySet().forEach(partition -> { - final int groupId = getTaskGroupIdForPartition(partition); - killTaskGroupForPartitions(ImmutableSet.of(partition), "DataSourceMetadata is updated while reset"); - taskGroups.remove(groupId); - partitionGroups.get(groupId).replaceAll((partitionId, offset) -> NOT_SET); - }); - } else { - throw new ISE("Unable to reset metadata"); - } - } else { - log.warn( - "Reset metadata topic [%s] and supervisor's topic [%s] do not match", - resetKafkaMetadata.getKafkaPartitions().getTopic(), - ioConfig.getTopic() - ); - } - } - } - - private void killTaskGroupForPartitions(Set partitions, String reasonFormat, Object... args) - { - for (Integer partition : partitions) { - killTasksInGroup(taskGroups.get(getTaskGroupIdForPartition(partition)), reasonFormat, args); - } - } - - private void killTasksInGroup(TaskGroup taskGroup, String reasonFormat, Object... args) - { - if (taskGroup != null) { - for (String taskId : taskGroup.tasks.keySet()) { - killTask(taskId, reasonFormat, args); - } - } - } - - @VisibleForTesting - void gracefulShutdownInternal() throws ExecutionException, InterruptedException, TimeoutException - { - // Prepare for shutdown by 1) killing all tasks that haven't been assigned to a worker yet, and 2) causing all - // running tasks to begin publishing by setting their startTime to a very long time ago so that the logic in - // checkTaskDuration() will be triggered. This is better than just telling these tasks to publish whatever they - // have, as replicas that are supposed to publish the same segment may not have read the same set of offsets. - for (TaskGroup taskGroup : taskGroups.values()) { - for (Entry entry : taskGroup.tasks.entrySet()) { - if (taskInfoProvider.getTaskLocation(entry.getKey()).equals(TaskLocation.unknown())) { - killTask(entry.getKey(), "Killing task for graceful shutdown"); - } else { - entry.getValue().startTime = DateTimes.EPOCH; - } - } - } - - checkTaskDuration(); - } - - @VisibleForTesting - void runInternal() throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException - { - possiblyRegisterListener(); - updatePartitionDataFromKafka(); - discoverTasks(); - updateTaskStatus(); - checkTaskDuration(); - checkPendingCompletionTasks(); - checkCurrentTaskState(); - - // if supervisor is not suspended, ensure required tasks are running - // if suspended, ensure tasks have been requested to gracefully stop - if (!spec.isSuspended()) { - log.info("[%s] supervisor is running.", dataSource); - createNewTasks(); - } else { - log.info("[%s] supervisor is suspended.", dataSource); - gracefulShutdownInternal(); - } - - if (log.isDebugEnabled()) { - log.debug(generateReport(true).toString()); - } else { - log.info(generateReport(false).toString()); - } - } - - String generateSequenceName( - Map startPartitions, - Optional minimumMessageTime, - Optional maximumMessageTime - ) - { - StringBuilder sb = new StringBuilder(); - - for (Entry entry : startPartitions.entrySet()) { - sb.append(StringUtils.format("+%d(%d)", entry.getKey(), entry.getValue())); - } - String partitionOffsetStr = sb.toString().substring(1); - - String minMsgTimeStr = (minimumMessageTime.isPresent() ? String.valueOf(minimumMessageTime.get().getMillis()) : ""); - String maxMsgTimeStr = (maximumMessageTime.isPresent() ? String.valueOf(maximumMessageTime.get().getMillis()) : ""); - - String dataSchema, tuningConfig; - try { - dataSchema = sortingMapper.writeValueAsString(spec.getDataSchema()); - tuningConfig = sortingMapper.writeValueAsString(taskTuningConfig); - } - catch (JsonProcessingException e) { - throw Throwables.propagate(e); - } - - String hashCode = DigestUtils.sha1Hex(dataSchema - + tuningConfig - + partitionOffsetStr - + minMsgTimeStr - + maxMsgTimeStr) - .substring(0, 15); - - return Joiner.on("_").join("index_kafka", dataSource, hashCode); - } - - @VisibleForTesting - protected void tryInit() - { - synchronized (stateChangeLock) { - if (started) { - log.warn("SUpervisor was already started, skipping init"); - return; - } - - if (stopped) { - log.warn("Supervisor was already stopped, skipping init."); - return; - } - - try { - consumer = getKafkaConsumer(); - - exec.submit( - () -> { - try { - long pollTimeout = Math.max(ioConfig.getPeriod().getMillis(), MAX_RUN_FREQUENCY_MILLIS); - while (!Thread.currentThread().isInterrupted() && !stopped) { - final Notice notice = notices.poll(pollTimeout, TimeUnit.MILLISECONDS); - if (notice == null) { - continue; - } - - try { - notice.handle(); - } - catch (Throwable e) { - log.makeAlert(e, "KafkaSupervisor[%s] failed to handle notice", dataSource) - .addData("noticeClass", notice.getClass().getSimpleName()) - .emit(); - } - } - } - catch (InterruptedException e) { - log.info("KafkaSupervisor[%s] interrupted, exiting", dataSource); - } - } - ); - firstRunTime = DateTimes.nowUtc().plus(ioConfig.getStartDelay()); - scheduledExec.scheduleAtFixedRate( - buildRunTask(), - ioConfig.getStartDelay().getMillis(), - Math.max(ioConfig.getPeriod().getMillis(), MAX_RUN_FREQUENCY_MILLIS), - TimeUnit.MILLISECONDS - ); - - reportingExec.scheduleAtFixedRate( - updateCurrentAndLatestOffsets(), - ioConfig.getStartDelay().getMillis() + INITIAL_GET_OFFSET_DELAY_MILLIS, // wait for tasks to start up - Math.max( - tuningConfig.getOffsetFetchPeriod().getMillis(), MINIMUM_GET_OFFSET_PERIOD_MILLIS - ), - TimeUnit.MILLISECONDS - ); - - reportingExec.scheduleAtFixedRate( - emitLag(), - ioConfig.getStartDelay().getMillis() + INITIAL_EMIT_LAG_METRIC_DELAY_MILLIS, // wait for tasks to start up - monitorSchedulerConfig.getEmitterPeriod().getMillis(), - TimeUnit.MILLISECONDS - ); - - started = true; - log.info( - "Started KafkaSupervisor[%s], first run in [%s], with spec: [%s]", - dataSource, - ioConfig.getStartDelay(), - spec.toString() - ); - } - catch (Exception e) { - if (consumer != null) { - consumer.close(); - } - initRetryCounter++; - log.makeAlert(e, "Exception starting KafkaSupervisor[%s]", dataSource) - .emit(); - - throw new RuntimeException(e); - } - } - } - - private KafkaConsumer getKafkaConsumer() - { - final Properties props = new Properties(); - - props.setProperty("metadata.max.age.ms", "10000"); - props.setProperty("group.id", StringUtils.format("kafka-supervisor-%s", RealtimeIndexTask.makeRandomId())); - - KafkaIndexTask.addConsumerPropertiesFromConfig(props, sortingMapper, ioConfig.getConsumerProperties()); - - props.setProperty("enable.auto.commit", "false"); - - ClassLoader currCtxCl = Thread.currentThread().getContextClassLoader(); - try { - Thread.currentThread().setContextClassLoader(getClass().getClassLoader()); - return new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer()); - } - finally { - Thread.currentThread().setContextClassLoader(currCtxCl); - } - } - - private void updatePartitionDataFromKafka() - { - List partitions; - try { - synchronized (consumerLock) { - partitions = consumer.partitionsFor(ioConfig.getTopic()); - } - } - catch (Exception e) { - log.warn( - e, - "Unable to get partition data from Kafka for brokers [%s], are the brokers up?", - ioConfig.getConsumerProperties().get(KafkaSupervisorIOConfig.BOOTSTRAP_SERVERS_KEY) - ); - return; - } - - int numPartitions = (partitions != null ? partitions.size() : 0); - - log.debug("Found [%d] Kafka partitions for topic [%s]", numPartitions, ioConfig.getTopic()); - - for (int partition = 0; partition < numPartitions; partition++) { - int taskGroupId = getTaskGroupIdForPartition(partition); - - ConcurrentHashMap partitionMap = partitionGroups.computeIfAbsent( - taskGroupId, - k -> new ConcurrentHashMap<>() - ); - - // The starting offset for a new partition in [partitionGroups] is initially set to NOT_SET; when a new task group - // is created and is assigned partitions, if the offset in [partitionGroups] is NOT_SET it will take the starting - // offset value from the metadata store, and if it can't find it there, from Kafka. Once a task begins - // publishing, the offset in partitionGroups will be updated to the ending offset of the publishing-but-not-yet- - // completed task, which will cause the next set of tasks to begin reading from where the previous task left - // off. If that previous task now fails, we will set the offset in [partitionGroups] back to NOT_SET which will - // cause successive tasks to again grab their starting offset from metadata store. This mechanism allows us to - // start up successive tasks without waiting for the previous tasks to succeed and still be able to handle task - // failures during publishing. - if (partitionMap.putIfAbsent(partition, NOT_SET) == null) { - log.info( - "New partition [%d] discovered for topic [%s], added to task group [%d]", - partition, - ioConfig.getTopic(), - taskGroupId - ); - } - } - } - - private void discoverTasks() throws ExecutionException, InterruptedException, TimeoutException - { - int taskCount = 0; - List futureTaskIds = new ArrayList<>(); - List> futures = new ArrayList<>(); - List tasks = taskStorage.getActiveTasks(); - final Map taskGroupsToVerify = new HashMap<>(); - - for (Task task : tasks) { - if (!(task instanceof KafkaIndexTask) || !dataSource.equals(task.getDataSource())) { - continue; - } - - taskCount++; - final KafkaIndexTask kafkaTask = (KafkaIndexTask) task; - final String taskId = task.getId(); - - // Determine which task group this task belongs to based on one of the partitions handled by this task. If we - // later determine that this task is actively reading, we will make sure that it matches our current partition - // allocation (getTaskGroupIdForPartition(partition) should return the same value for every partition being read - // by this task) and kill it if it is not compatible. If the task is instead found to be in the publishing - // state, we will permit it to complete even if it doesn't match our current partition allocation to support - // seamless schema migration. - - Iterator it = kafkaTask.getIOConfig().getStartPartitions().getPartitionOffsetMap().keySet().iterator(); - final Integer taskGroupId = (it.hasNext() ? getTaskGroupIdForPartition(it.next()) : null); - - if (taskGroupId != null) { - // check to see if we already know about this task, either in [taskGroups] or in [pendingCompletionTaskGroups] - // and if not add it to taskGroups or pendingCompletionTaskGroups (if status = PUBLISHING) - TaskGroup taskGroup = taskGroups.get(taskGroupId); - if (!isTaskInPendingCompletionGroups(taskId) && (taskGroup == null || !taskGroup.tasks.containsKey(taskId))) { - - futureTaskIds.add(taskId); - futures.add( - Futures.transform( - taskClient.getStatusAsync(taskId), new Function() - { - @Override - public Boolean apply(KafkaIndexTask.Status status) - { - try { - log.debug("Task [%s], status [%s]", taskId, status); - if (status == KafkaIndexTask.Status.PUBLISHING) { - kafkaTask.getIOConfig().getStartPartitions().getPartitionOffsetMap().keySet().forEach( - partition -> addDiscoveredTaskToPendingCompletionTaskGroups( - getTaskGroupIdForPartition(partition), - taskId, - kafkaTask.getIOConfig() - .getStartPartitions() - .getPartitionOffsetMap() - ) - ); - - // update partitionGroups with the publishing task's offsets (if they are greater than what is - // existing) so that the next tasks will start reading from where this task left off - Map publishingTaskEndOffsets = taskClient.getEndOffsets(taskId); - - for (Entry entry : publishingTaskEndOffsets.entrySet()) { - Integer partition = entry.getKey(); - Long offset = entry.getValue(); - ConcurrentHashMap partitionOffsets = partitionGroups.get( - getTaskGroupIdForPartition(partition) - ); - - boolean succeeded; - do { - succeeded = true; - Long previousOffset = partitionOffsets.putIfAbsent(partition, offset); - if (previousOffset != null && previousOffset < offset) { - succeeded = partitionOffsets.replace(partition, previousOffset, offset); - } - } while (!succeeded); - } - } else { - for (Integer partition : kafkaTask.getIOConfig() - .getStartPartitions() - .getPartitionOffsetMap() - .keySet()) { - if (!taskGroupId.equals(getTaskGroupIdForPartition(partition))) { - log.warn( - "Stopping task [%s] which does not match the expected partition allocation", - taskId - ); - try { - stopTask(taskId, false).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - catch (InterruptedException | ExecutionException | TimeoutException e) { - log.warn(e, "Exception while stopping task"); - } - return false; - } - } - // make sure the task's io and tuning configs match with the supervisor config - // if it is current then only create corresponding taskGroup if it does not exist - if (!isTaskCurrent(taskGroupId, taskId)) { - log.info( - "Stopping task [%s] which does not match the expected parameters and ingestion spec", - taskId - ); - try { - stopTask(taskId, false).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - catch (InterruptedException | ExecutionException | TimeoutException e) { - log.warn(e, "Exception while stopping task"); - } - return false; - } else { - final TaskGroup taskGroup = taskGroups.computeIfAbsent( - taskGroupId, - k -> { - log.info("Creating a new task group for taskGroupId[%d]", taskGroupId); - return new TaskGroup( - taskGroupId, - ImmutableMap.copyOf( - kafkaTask.getIOConfig().getStartPartitions().getPartitionOffsetMap() - ), - kafkaTask.getIOConfig().getMinimumMessageTime(), - kafkaTask.getIOConfig().getMaximumMessageTime() - ); - } - ); - taskGroupsToVerify.put(taskGroupId, taskGroup); - final TaskData prevTaskData = taskGroup.tasks.putIfAbsent(taskId, new TaskData()); - if (prevTaskData != null) { - throw new ISE( - "WTH? a taskData[%s] already exists for new task[%s]", - prevTaskData, - taskId - ); - } - } - } - return true; - } - catch (Throwable t) { - log.error(t, "Something bad while discovering task [%s]", taskId); - return null; - } - } - }, workerExec - ) - ); - } - } - } - - List results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - for (int i = 0; i < results.size(); i++) { - if (results.get(i) == null) { - String taskId = futureTaskIds.get(i); - killTask(taskId, "Task [%s] failed to return status, killing task", taskId); - } - } - log.debug("Found [%d] Kafka indexing tasks for dataSource [%s]", taskCount, dataSource); - - // make sure the checkpoints are consistent with each other and with the metadata store - verifyAndMergeCheckpoints(taskGroupsToVerify.values()); - } - - private void verifyAndMergeCheckpoints(final Collection taskGroupsToVerify) - { - final List> futures = new ArrayList<>(); - for (TaskGroup taskGroup : taskGroupsToVerify) { - futures.add(workerExec.submit(() -> verifyAndMergeCheckpoints(taskGroup))); - } - try { - Futures.allAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - catch (InterruptedException | ExecutionException | TimeoutException e) { - throw new RuntimeException(e); - } - } - - /** - * This method does two things - - * 1. Makes sure the checkpoints information in the taskGroup is consistent with that of the tasks, if not kill - * inconsistent tasks. - * 2. truncates the checkpoints in the taskGroup corresponding to which segments have been published, so that any newly - * created tasks for the taskGroup start indexing from after the latest published offsets. - */ - private void verifyAndMergeCheckpoints(final TaskGroup taskGroup) - { - final int groupId = taskGroup.groupId; - final List>>> taskSequences = new ArrayList<>(); - final List>>> futures = new ArrayList<>(); - final List taskIds = new ArrayList<>(); - - for (String taskId : taskGroup.taskIds()) { - final ListenableFuture>> checkpointsFuture = taskClient.getCheckpointsAsync( - taskId, - true - ); - taskIds.add(taskId); - futures.add(checkpointsFuture); - } - - try { - List>> futuresResult = Futures.successfulAsList(futures) - .get(futureTimeoutInSeconds, TimeUnit.SECONDS); - - for (int i = 0; i < futuresResult.size(); i++) { - final TreeMap> checkpoints = futuresResult.get(i); - final String taskId = taskIds.get(i); - if (checkpoints == null) { - try { - // catch the exception in failed futures - futures.get(i).get(); - } - catch (Exception e) { - log.error(e, "Problem while getting checkpoints for task [%s], killing the task", taskId); - killTask(taskId, "Exception[%s] while getting checkpoints", e.getClass()); - taskGroup.tasks.remove(taskId); - } - } else if (checkpoints.isEmpty()) { - log.warn("Ignoring task [%s], as probably it is not started running yet", taskId); - } else { - taskSequences.add(new Pair<>(taskId, checkpoints)); - } - } - } - catch (Exception e) { - throw new RuntimeException(e); - } - - final KafkaDataSourceMetadata latestDataSourceMetadata = (KafkaDataSourceMetadata) indexerMetadataStorageCoordinator - .getDataSourceMetadata(dataSource); - final boolean hasValidOffsetsFromDb = latestDataSourceMetadata != null && - latestDataSourceMetadata.getKafkaPartitions() != null && - ioConfig.getTopic().equals( - latestDataSourceMetadata.getKafkaPartitions().getTopic() - ); - final Map latestOffsetsFromDb; - if (hasValidOffsetsFromDb) { - latestOffsetsFromDb = latestDataSourceMetadata.getKafkaPartitions().getPartitionOffsetMap(); - } else { - latestOffsetsFromDb = null; - } - - // order tasks of this taskGroup by the latest sequenceId - taskSequences.sort((o1, o2) -> o2.rhs.firstKey().compareTo(o1.rhs.firstKey())); - - final Set tasksToKill = new HashSet<>(); - final AtomicInteger earliestConsistentSequenceId = new AtomicInteger(-1); - int taskIndex = 0; - - while (taskIndex < taskSequences.size()) { - TreeMap> taskCheckpoints = taskSequences.get(taskIndex).rhs; - String taskId = taskSequences.get(taskIndex).lhs; - if (earliestConsistentSequenceId.get() == -1) { - // find the first replica task with earliest sequenceId consistent with datasource metadata in the metadata - // store - if (taskCheckpoints.entrySet().stream().anyMatch( - sequenceCheckpoint -> sequenceCheckpoint.getValue().entrySet().stream().allMatch( - partitionOffset -> Longs.compare( - partitionOffset.getValue(), - latestOffsetsFromDb == null ? - partitionOffset.getValue() : - latestOffsetsFromDb.getOrDefault(partitionOffset.getKey(), partitionOffset.getValue()) - ) == 0) && earliestConsistentSequenceId.compareAndSet(-1, sequenceCheckpoint.getKey())) || ( - someTaskGroupsPendingCompletion(groupId) - && earliestConsistentSequenceId.compareAndSet(-1, taskCheckpoints.firstKey()))) { - final SortedMap> latestCheckpoints = new TreeMap<>( - taskCheckpoints.tailMap(earliestConsistentSequenceId.get()) - ); - log.info("Setting taskGroup sequences to [%s] for group [%d]", latestCheckpoints, groupId); - taskGroup.sequenceOffsets.clear(); - taskGroup.sequenceOffsets.putAll(latestCheckpoints); - } else { - log.debug( - "Adding task [%s] to kill list, checkpoints[%s], latestoffsets from DB [%s]", - taskId, - taskCheckpoints, - latestOffsetsFromDb - ); - tasksToKill.add(taskId); - } - } else { - // check consistency with taskGroup sequences - if (taskCheckpoints.get(taskGroup.sequenceOffsets.firstKey()) == null - || !(taskCheckpoints.get(taskGroup.sequenceOffsets.firstKey()) - .equals(taskGroup.sequenceOffsets.firstEntry().getValue())) - || taskCheckpoints.tailMap(taskGroup.sequenceOffsets.firstKey()).size() - != taskGroup.sequenceOffsets.size()) { - log.debug( - "Adding task [%s] to kill list, checkpoints[%s], taskgroup checkpoints [%s]", - taskId, - taskCheckpoints, - taskGroup.sequenceOffsets - ); - tasksToKill.add(taskId); - } - } - taskIndex++; - } - - if ((tasksToKill.size() > 0 && tasksToKill.size() == taskGroup.tasks.size()) || - (taskGroup.tasks.size() == 0 && !someTaskGroupsPendingCompletion(groupId))) { - // killing all tasks or no task left in the group ? - // clear state about the taskgroup so that get latest offset information is fetched from metadata store - log.warn("Clearing task group [%d] information as no valid tasks left the group", groupId); - taskGroups.remove(groupId); - partitionGroups.get(groupId).replaceAll((partition, offset) -> NOT_SET); - } - - taskSequences.stream().filter(taskIdSequences -> tasksToKill.contains(taskIdSequences.lhs)).forEach( - sequenceCheckpoint -> { - killTask( - sequenceCheckpoint.lhs, - "Killing task [%s], as its checkpoints [%s] are not consistent with group checkpoints[%s] or latest " - + "persisted offsets in metadata store [%s]", - sequenceCheckpoint.lhs, - sequenceCheckpoint.rhs, - taskGroup.sequenceOffsets, - latestOffsetsFromDb - ); - taskGroup.tasks.remove(sequenceCheckpoint.lhs); - } - ); - } - - private void addDiscoveredTaskToPendingCompletionTaskGroups( - int groupId, - String taskId, - Map startingPartitions - ) - { - final CopyOnWriteArrayList taskGroupList = pendingCompletionTaskGroups.computeIfAbsent( - groupId, - k -> new CopyOnWriteArrayList<>() - ); - for (TaskGroup taskGroup : taskGroupList) { - if (taskGroup.partitionOffsets.equals(startingPartitions)) { - if (taskGroup.tasks.putIfAbsent(taskId, new TaskData()) == null) { - log.info("Added discovered task [%s] to existing pending task group [%s]", taskId, groupId); - } - return; - } - } - - log.info("Creating new pending completion task group [%s] for discovered task [%s]", groupId, taskId); - - // reading the minimumMessageTime & maximumMessageTime from the publishing task and setting it here is not necessary as this task cannot - // change to a state where it will read any more events - TaskGroup newTaskGroup = new TaskGroup( - groupId, - ImmutableMap.copyOf(startingPartitions), - Optional.absent(), - Optional.absent() - ); - - newTaskGroup.tasks.put(taskId, new TaskData()); - newTaskGroup.completionTimeout = DateTimes.nowUtc().plus(ioConfig.getCompletionTimeout()); - - taskGroupList.add(newTaskGroup); - } - - private void updateTaskStatus() throws ExecutionException, InterruptedException, TimeoutException - { - final List> futures = new ArrayList<>(); - final List futureTaskIds = new ArrayList<>(); - - // update status (and startTime if unknown) of current tasks in taskGroups - for (TaskGroup group : taskGroups.values()) { - for (Entry entry : group.tasks.entrySet()) { - final String taskId = entry.getKey(); - final TaskData taskData = entry.getValue(); - - if (taskData.startTime == null) { - futureTaskIds.add(taskId); - futures.add( - Futures.transform( - taskClient.getStartTimeAsync(taskId), new Function() - { - @Nullable - @Override - public Boolean apply(@Nullable DateTime startTime) - { - if (startTime == null) { - return false; - } - - taskData.startTime = startTime; - long millisRemaining = ioConfig.getTaskDuration().getMillis() - - (System.currentTimeMillis() - taskData.startTime.getMillis()); - if (millisRemaining > 0) { - scheduledExec.schedule( - buildRunTask(), - millisRemaining + MAX_RUN_FREQUENCY_MILLIS, - TimeUnit.MILLISECONDS - ); - } - - return true; - } - }, workerExec - ) - ); - } - - taskData.status = taskStorage.getStatus(taskId).get(); - } - } - - // update status of pending completion tasks in pendingCompletionTaskGroups - for (List taskGroups : pendingCompletionTaskGroups.values()) { - for (TaskGroup group : taskGroups) { - for (Entry entry : group.tasks.entrySet()) { - entry.getValue().status = taskStorage.getStatus(entry.getKey()).get(); - } - } - } - - List results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - for (int i = 0; i < results.size(); i++) { - // false means the task hasn't started running yet and that's okay; null means it should be running but the HTTP - // request threw an exception so kill the task - if (results.get(i) == null) { - String taskId = futureTaskIds.get(i); - killTask(taskId, "Task [%s] failed to return start time, killing task", taskId); - } - } - } - - private void checkTaskDuration() throws InterruptedException, ExecutionException, TimeoutException - { - final List>> futures = new ArrayList<>(); - final List futureGroupIds = new ArrayList<>(); - - for (Entry entry : taskGroups.entrySet()) { - Integer groupId = entry.getKey(); - TaskGroup group = entry.getValue(); - - // find the longest running task from this group - DateTime earliestTaskStart = DateTimes.nowUtc(); - for (TaskData taskData : group.tasks.values()) { - // startTime can be null if kafkaSupervisor is stopped gracefully before processing any runNotice - if (taskData.startTime != null && earliestTaskStart.isAfter(taskData.startTime)) { - earliestTaskStart = taskData.startTime; - } - } - - // if this task has run longer than the configured duration, signal all tasks in the group to persist - if (earliestTaskStart.plus(ioConfig.getTaskDuration()).isBeforeNow()) { - log.info("Task group [%d] has run for [%s]", groupId, ioConfig.getTaskDuration()); - futureGroupIds.add(groupId); - futures.add(checkpointTaskGroup(group, true)); - } - } - - List> results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - for (int j = 0; j < results.size(); j++) { - Integer groupId = futureGroupIds.get(j); - TaskGroup group = taskGroups.get(groupId); - Map endOffsets = results.get(j); - - if (endOffsets != null) { - // set a timeout and put this group in pendingCompletionTaskGroups so that it can be monitored for completion - group.completionTimeout = DateTimes.nowUtc().plus(ioConfig.getCompletionTimeout()); - pendingCompletionTaskGroups.computeIfAbsent(groupId, k -> new CopyOnWriteArrayList<>()).add(group); - - // set endOffsets as the next startOffsets - for (Entry entry : endOffsets.entrySet()) { - partitionGroups.get(groupId).put(entry.getKey(), entry.getValue()); - } - } else { - for (String id : group.taskIds()) { - killTask( - id, - "All tasks in group [%s] failed to transition to publishing state", - groupId - ); - } - // clear partitionGroups, so that latest offsets from db is used as start offsets not the stale ones - // if tasks did some successful incremental handoffs - partitionGroups.get(groupId).replaceAll((partition, offset) -> NOT_SET); - } - - // remove this task group from the list of current task groups now that it has been handled - taskGroups.remove(groupId); - } - } - - private ListenableFuture> checkpointTaskGroup(final TaskGroup taskGroup, final boolean finalize) - { - if (finalize) { - // 1) Check if any task completed (in which case we're done) and kill unassigned tasks - Iterator> i = taskGroup.tasks.entrySet().iterator(); - while (i.hasNext()) { - Entry taskEntry = i.next(); - String taskId = taskEntry.getKey(); - TaskData task = taskEntry.getValue(); - - // task.status can be null if kafkaSupervisor is stopped gracefully before processing any runNotice. - if (task.status != null) { - if (task.status.isSuccess()) { - // If any task in this group has already completed, stop the rest of the tasks in the group and return. - // This will cause us to create a new set of tasks next cycle that will start from the offsets in - // metadata store (which will have advanced if we succeeded in publishing and will remain the same if - // publishing failed and we need to re-ingest) - return Futures.transform( - stopTasksInGroup(taskGroup, "task[%s] succeeded in the taskGroup", task.status.getId()), - new Function>() - { - @Nullable - @Override - public Map apply(@Nullable Object input) - { - return null; - } - } - ); - } - - if (task.status.isRunnable()) { - if (taskInfoProvider.getTaskLocation(taskId).equals(TaskLocation.unknown())) { - killTask(taskId, "Killing task [%s] which hasn't been assigned to a worker", taskId); - i.remove(); - } - } - } - } - } - - // 2) Pause running tasks - final List>> pauseFutures = new ArrayList<>(); - final List pauseTaskIds = ImmutableList.copyOf(taskGroup.taskIds()); - for (final String taskId : pauseTaskIds) { - pauseFutures.add(taskClient.pauseAsync(taskId)); - } - - return Futures.transform( - Futures.successfulAsList(pauseFutures), new Function>, Map>() - { - @Nullable - @Override - public Map apply(List> input) - { - // 3) Build a map of the highest offset read by any task in the group for each partition - final Map endOffsets = new HashMap<>(); - for (int i = 0; i < input.size(); i++) { - final Map result = input.get(i); - final String taskId = pauseTaskIds.get(i); - - if (result == null) { - // Get the exception - final Throwable pauseException; - try { - // The below get should throw ExecutionException since result is null. - final Map pauseResult = pauseFutures.get(i).get(); - throw new ISE( - "WTH? The pause request for task [%s] is supposed to fail, but returned [%s]", - taskId, - pauseResult - ); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - catch (ExecutionException e) { - pauseException = e.getCause() == null ? e : e.getCause(); - } - - killTask( - taskId, - "An exception occured while waiting for task [%s] to pause: [%s]", - taskId, - pauseException - ); - taskGroup.tasks.remove(taskId); - - } else if (result.isEmpty()) { - killTask(taskId, "Task [%s] returned empty offsets after pause", taskId); - taskGroup.tasks.remove(taskId); - } else { // otherwise build a map of the highest offsets seen - for (Entry offset : result.entrySet()) { - if (!endOffsets.containsKey(offset.getKey()) - || endOffsets.get(offset.getKey()).compareTo(offset.getValue()) < 0) { - endOffsets.put(offset.getKey(), offset.getValue()); - } - } - } - } - - // 4) Set the end offsets for each task to the values from step 3 and resume the tasks. All the tasks should - // finish reading and start publishing within a short period, depending on how in sync the tasks were. - final List> setEndOffsetFutures = new ArrayList<>(); - final List setEndOffsetTaskIds = ImmutableList.copyOf(taskGroup.taskIds()); - - if (setEndOffsetTaskIds.isEmpty()) { - log.info("All tasks in taskGroup [%d] have failed, tasks will be re-created", taskGroup.groupId); - return null; - } - - try { - - if (endOffsets.equals(taskGroup.sequenceOffsets.lastEntry().getValue())) { - log.warn( - "Checkpoint [%s] is same as the start offsets [%s] of latest sequence for the task group [%d]", - endOffsets, - taskGroup.sequenceOffsets.lastEntry().getValue(), - taskGroup.groupId - ); - } - - log.info( - "Setting endOffsets for tasks in taskGroup [%d] to %s and resuming", - taskGroup.groupId, - endOffsets - ); - for (final String taskId : setEndOffsetTaskIds) { - setEndOffsetFutures.add(taskClient.setEndOffsetsAsync(taskId, endOffsets, finalize)); - } - - List results = Futures.successfulAsList(setEndOffsetFutures) - .get(futureTimeoutInSeconds, TimeUnit.SECONDS); - for (int i = 0; i < results.size(); i++) { - if (results.get(i) == null || !results.get(i)) { - String taskId = setEndOffsetTaskIds.get(i); - killTask( - taskId, - "Task [%s] failed to respond to [set end offsets] in a timely manner, killing task", - taskId - ); - taskGroup.tasks.remove(taskId); - } - } - } - catch (Exception e) { - log.error("Something bad happened [%s]", e.getMessage()); - Throwables.propagate(e); - } - - if (taskGroup.tasks.isEmpty()) { - log.info("All tasks in taskGroup [%d] have failed, tasks will be re-created", taskGroup.groupId); - return null; - } - - return endOffsets; - } - }, workerExec - ); - } - - /** - * Monitors [pendingCompletionTaskGroups] for tasks that have completed. If any task in a task group has completed, we - * can safely stop the rest of the tasks in that group. If a task group has exceeded its publishing timeout, then - * we need to stop all tasks in not only that task group but also 1) any subsequent task group that is also pending - * completion and 2) the current task group that is running, because the assumption that we have handled up to the - * starting offset for subsequent task groups is no longer valid, and subsequent tasks would fail as soon as they - * attempted to publish because of the contiguous range consistency check. - */ - private void checkPendingCompletionTasks() throws ExecutionException, InterruptedException, TimeoutException - { - List> futures = new ArrayList<>(); - - for (Entry> pendingGroupList : pendingCompletionTaskGroups.entrySet()) { - - boolean stopTasksInTaskGroup = false; - Integer groupId = pendingGroupList.getKey(); - CopyOnWriteArrayList taskGroupList = pendingGroupList.getValue(); - List toRemove = new ArrayList<>(); - - for (TaskGroup group : taskGroupList) { - boolean foundSuccess = false, entireTaskGroupFailed = false; - - if (stopTasksInTaskGroup) { - // One of the earlier groups that was handling the same partition set timed out before the segments were - // published so stop any additional groups handling the same partition set that are pending completion. - futures.add( - stopTasksInGroup( - group, - "one of earlier groups that was handling the same partition set timed out before publishing segments" - ) - ); - toRemove.add(group); - continue; - } - - Iterator> iTask = group.tasks.entrySet().iterator(); - while (iTask.hasNext()) { - final Entry entry = iTask.next(); - final String taskId = entry.getKey(); - final TaskData taskData = entry.getValue(); - - Preconditions.checkNotNull(taskData.status, "WTH? task[%s] has a null status", taskId); - - if (taskData.status.isFailure()) { - iTask.remove(); // remove failed task - if (group.tasks.isEmpty()) { - // if all tasks in the group have failed, just nuke all task groups with this partition set and restart - entireTaskGroupFailed = true; - break; - } - } - - if (taskData.status.isSuccess()) { - // If one of the pending completion tasks was successful, stop the rest of the tasks in the group as - // we no longer need them to publish their segment. - futures.add( - stopTasksInGroup(group, "Task [%s] completed successfully, stopping tasks %s", taskId, group.taskIds()) - ); - foundSuccess = true; - toRemove.add(group); // remove the TaskGroup from the list of pending completion task groups - break; // skip iterating the rest of the tasks in this group as they've all been stopped now - } - } - - if ((!foundSuccess && group.completionTimeout.isBeforeNow()) || entireTaskGroupFailed) { - if (entireTaskGroupFailed) { - log.warn("All tasks in group [%d] failed to publish, killing all tasks for these partitions", groupId); - } else { - log.makeAlert( - "No task in [%s] for taskGroup [%d] succeeded before the completion timeout elapsed [%s]!", - group.taskIds(), - groupId, - ioConfig.getCompletionTimeout() - ).emit(); - } - - // reset partitions offsets for this task group so that they will be re-read from metadata storage - partitionGroups.get(groupId).replaceAll((partition, offset) -> NOT_SET); - // kill all the tasks in this pending completion group - killTasksInGroup( - group, - "No task in pending completion taskGroup[%d] succeeded before completion timeout elapsed", - groupId - ); - // set a flag so the other pending completion groups for this set of partitions will also stop - stopTasksInTaskGroup = true; - - // kill all the tasks in the currently reading task group and remove the bad task group - killTasksInGroup( - taskGroups.remove(groupId), - "No task in the corresponding pending completion taskGroup[%d] succeeded before completion timeout elapsed", - groupId - ); - toRemove.add(group); - } - } - - taskGroupList.removeAll(toRemove); - } - - // wait for all task shutdowns to complete before returning - Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - - private void checkCurrentTaskState() throws ExecutionException, InterruptedException, TimeoutException - { - List> futures = new ArrayList<>(); - Iterator> iTaskGroups = taskGroups.entrySet().iterator(); - while (iTaskGroups.hasNext()) { - Entry taskGroupEntry = iTaskGroups.next(); - Integer groupId = taskGroupEntry.getKey(); - TaskGroup taskGroup = taskGroupEntry.getValue(); - - // Iterate the list of known tasks in this group and: - // 1) Kill any tasks which are not "current" (have the partitions, starting offsets, and minimumMessageTime - // & maximumMessageTime (if applicable) in [taskGroups]) - // 2) Remove any tasks that have failed from the list - // 3) If any task completed successfully, stop all the tasks in this group and move to the next group - - log.debug("Task group [%d] pre-pruning: %s", groupId, taskGroup.taskIds()); - - Iterator> iTasks = taskGroup.tasks.entrySet().iterator(); - while (iTasks.hasNext()) { - Entry task = iTasks.next(); - String taskId = task.getKey(); - TaskData taskData = task.getValue(); - - // stop and remove bad tasks from the task group - if (!isTaskCurrent(groupId, taskId)) { - log.info("Stopping task [%s] which does not match the expected offset range and ingestion spec", taskId); - futures.add(stopTask(taskId, false)); - iTasks.remove(); - continue; - } - - Preconditions.checkNotNull(taskData.status, "WTH? task[%s] has a null status", taskId); - - // remove failed tasks - if (taskData.status.isFailure()) { - iTasks.remove(); - continue; - } - - // check for successful tasks, and if we find one, stop all tasks in the group and remove the group so it can - // be recreated with the next set of offsets - if (taskData.status.isSuccess()) { - futures.add(stopTasksInGroup(taskGroup, "task[%s] succeeded in the same taskGroup", taskData.status.getId())); - iTaskGroups.remove(); - break; - } - } - log.debug("Task group [%d] post-pruning: %s", groupId, taskGroup.taskIds()); - } - - // wait for all task shutdowns to complete before returning - Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - - void createNewTasks() throws JsonProcessingException - { - // update the checkpoints in the taskGroup to latest ones so that new tasks do not read what is already published - verifyAndMergeCheckpoints( - taskGroups.values() - .stream() - .filter(taskGroup -> taskGroup.tasks.size() < ioConfig.getReplicas()) - .collect(Collectors.toList()) - ); - - // check that there is a current task group for each group of partitions in [partitionGroups] - for (Integer groupId : partitionGroups.keySet()) { - if (!taskGroups.containsKey(groupId)) { - log.info("Creating new task group [%d] for partitions %s", groupId, partitionGroups.get(groupId).keySet()); - - Optional minimumMessageTime = (ioConfig.getLateMessageRejectionPeriod().isPresent() ? Optional.of( - DateTimes.nowUtc().minus(ioConfig.getLateMessageRejectionPeriod().get()) - ) : Optional.absent()); - - Optional maximumMessageTime = (ioConfig.getEarlyMessageRejectionPeriod().isPresent() ? Optional.of( - DateTimes.nowUtc().plus(ioConfig.getTaskDuration()).plus(ioConfig.getEarlyMessageRejectionPeriod().get()) - ) : Optional.absent()); - - final TaskGroup taskGroup = new TaskGroup( - groupId, - generateStartingOffsetsForPartitionGroup(groupId), - minimumMessageTime, - maximumMessageTime - ); - taskGroups.put( - groupId, - taskGroup - ); - } - } - - // iterate through all the current task groups and make sure each one has the desired number of replica tasks - boolean createdTask = false; - for (Entry entry : taskGroups.entrySet()) { - TaskGroup taskGroup = entry.getValue(); - Integer groupId = entry.getKey(); + private final ServiceEmitter emitter; + private final DruidMonitorSchedulerConfig monitorSchedulerConfig; + private volatile Map latestSequenceFromStream; - if (ioConfig.getReplicas() > taskGroup.tasks.size()) { - log.info( - "Number of tasks [%d] does not match configured numReplicas [%d] in task group [%d], creating more tasks", - taskGroup.tasks.size(), ioConfig.getReplicas(), groupId - ); - createKafkaTasksForGroup(groupId, ioConfig.getReplicas() - taskGroup.tasks.size()); - createdTask = true; - } - } - if (createdTask && firstRunTime.isBeforeNow()) { - // Schedule a run event after a short delay to update our internal data structures with the new tasks that were - // just created. This is mainly for the benefit of the status API in situations where the run period is lengthy. - scheduledExec.schedule(buildRunTask(), 5000, TimeUnit.MILLISECONDS); - } - } + private final KafkaSupervisorSpec spec; - private void createKafkaTasksForGroup(int groupId, int replicas) throws JsonProcessingException + public KafkaSupervisor( + final TaskStorage taskStorage, + final TaskMaster taskMaster, + final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + final KafkaIndexTaskClientFactory taskClientFactory, + final ObjectMapper mapper, + final KafkaSupervisorSpec spec, + final RowIngestionMetersFactory rowIngestionMetersFactory + ) { - Map startPartitions = taskGroups.get(groupId).partitionOffsets; - Map endPartitions = new HashMap<>(); - - for (Integer partition : startPartitions.keySet()) { - endPartitions.put(partition, Long.MAX_VALUE); - } - TaskGroup group = taskGroups.get(groupId); - - Map consumerProperties = new HashMap<>(ioConfig.getConsumerProperties()); - DateTime minimumMessageTime = taskGroups.get(groupId).minimumMessageTime.orNull(); - DateTime maximumMessageTime = taskGroups.get(groupId).maximumMessageTime.orNull(); - - KafkaIOConfig kafkaIOConfig = new KafkaIOConfig( - groupId, - group.baseSequenceName, - new KafkaPartitions(ioConfig.getTopic(), startPartitions), - new KafkaPartitions(ioConfig.getTopic(), endPartitions), - consumerProperties, - true, - minimumMessageTime, - maximumMessageTime, - ioConfig.isSkipOffsetGaps() + super( + StringUtils.format("KafkaSupervisor-%s", spec.getDataSchema().getDataSource()), + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + mapper, + spec, + rowIngestionMetersFactory, + false ); - final String checkpoints = sortingMapper.writerWithType(new TypeReference>>() - { - }).writeValueAsString(taskGroups.get(groupId).sequenceOffsets); - final Map context = spec.getContext() == null - ? ImmutableMap.of( - "checkpoints", - checkpoints, - IS_INCREMENTAL_HANDOFF_SUPPORTED, - true - ) - : ImmutableMap.builder() - .put("checkpoints", checkpoints) - .put(IS_INCREMENTAL_HANDOFF_SUPPORTED, true) - .putAll(spec.getContext()) - .build(); - for (int i = 0; i < replicas; i++) { - String taskId = Joiner.on("_").join(group.baseSequenceName, RealtimeIndexTask.makeRandomId()); - KafkaIndexTask indexTask = new KafkaIndexTask( - taskId, - new TaskResource(group.baseSequenceName, 1), - spec.getDataSchema(), - taskTuningConfig, - kafkaIOConfig, - context, - null, - null, - rowIngestionMetersFactory, - sortingMapper - ); - - Optional taskQueue = taskMaster.getTaskQueue(); - if (taskQueue.isPresent()) { - try { - taskQueue.get().add(indexTask); - } - catch (EntryExistsException e) { - log.error("Tried to add task [%s] but it already exists", indexTask.getId()); - } - } else { - log.error("Failed to get task queue because I'm not the leader!"); - } - } - } - - private ImmutableMap generateStartingOffsetsForPartitionGroup(int groupId) - { - ImmutableMap.Builder builder = ImmutableMap.builder(); - for (Entry entry : partitionGroups.get(groupId).entrySet()) { - Integer partition = entry.getKey(); - Long offset = entry.getValue(); - - if (offset != null && offset != NOT_SET) { - // if we are given a startingOffset (set by a previous task group which is pending completion) then use it - builder.put(partition, offset); - } else { - // if we don't have a startingOffset (first run or we had some previous failures and reset the offsets) then - // get the offset from metadata storage (if available) or Kafka (otherwise) - builder.put(partition, getOffsetFromStorageForPartition(partition)); - } - } - return builder.build(); - } - - /** - * Queries the dataSource metadata table to see if there is a previous ending offset for this partition. If it doesn't - * find any data, it will retrieve the latest or earliest Kafka offset depending on the useEarliestOffset config. - */ - private long getOffsetFromStorageForPartition(int partition) - { - long offset; - final Map metadataOffsets = getOffsetsFromMetadataStorage(); - if (metadataOffsets.get(partition) != null) { - offset = metadataOffsets.get(partition); - log.debug("Getting offset [%,d] from metadata storage for partition [%d]", offset, partition); - - long latestKafkaOffset = getOffsetFromKafkaForPartition(partition, false); - if (offset > latestKafkaOffset) { - throw new ISE( - "Offset in metadata storage [%,d] > latest Kafka offset [%,d] for partition[%d] dataSource[%s]. If these " - + "messages are no longer available (perhaps you deleted and re-created your Kafka topic) you can use the " - + "supervisor reset API to restart ingestion.", - offset, - latestKafkaOffset, - partition, - dataSource - ); - } - - } else { - offset = getOffsetFromKafkaForPartition(partition, ioConfig.isUseEarliestOffset()); - log.debug("Getting offset [%,d] from Kafka for partition [%d]", offset, partition); - } - - return offset; - } - - private Map getOffsetsFromMetadataStorage() - { - final DataSourceMetadata dataSourceMetadata = indexerMetadataStorageCoordinator.getDataSourceMetadata(dataSource); - if (dataSourceMetadata instanceof KafkaDataSourceMetadata) { - KafkaPartitions partitions = ((KafkaDataSourceMetadata) dataSourceMetadata).getKafkaPartitions(); - if (partitions != null) { - if (!ioConfig.getTopic().equals(partitions.getTopic())) { - log.warn( - "Topic in metadata storage [%s] doesn't match spec topic [%s], ignoring stored offsets", - partitions.getTopic(), - ioConfig.getTopic() - ); - return Collections.emptyMap(); - } else if (partitions.getPartitionOffsetMap() != null) { - return partitions.getPartitionOffsetMap(); - } - } - } - - return Collections.emptyMap(); + this.spec = spec; + this.emitter = spec.getEmitter(); + this.monitorSchedulerConfig = spec.getMonitorSchedulerConfig(); } - private long getOffsetFromKafkaForPartition(int partition, boolean useEarliestOffset) - { - synchronized (consumerLock) { - TopicPartition topicPartition = new TopicPartition(ioConfig.getTopic(), partition); - if (!consumer.assignment().contains(topicPartition)) { - consumer.assign(Collections.singletonList(topicPartition)); - } - - if (useEarliestOffset) { - consumer.seekToBeginning(Collections.singletonList(topicPartition)); - } else { - consumer.seekToEnd(Collections.singletonList(topicPartition)); - } - - return consumer.position(topicPartition); - } - } - /** - * Compares the sequence name from the task with one generated for the task's group ID and returns false if they do - * not match. The sequence name is generated from a hash of the dataSchema, tuningConfig, starting offsets, and the - * minimumMessageTime or maximumMessageTime if set. - */ - private boolean isTaskCurrent(int taskGroupId, String taskId) + @Override + protected RecordSupplier setupRecordSupplier() { - Optional taskOptional = taskStorage.getTask(taskId); - if (!taskOptional.isPresent() || !(taskOptional.get() instanceof KafkaIndexTask)) { - return false; - } - - String taskSequenceName = ((KafkaIndexTask) taskOptional.get()).getIOConfig().getBaseSequenceName(); - if (taskGroups.get(taskGroupId) != null) { - return Preconditions - .checkNotNull(taskGroups.get(taskGroupId), "null taskGroup for taskId[%s]", taskGroupId) - .baseSequenceName - .equals(taskSequenceName); - } else { - return generateSequenceName( - ((KafkaIndexTask) taskOptional.get()).getIOConfig() - .getStartPartitions() - .getPartitionOffsetMap(), - ((KafkaIndexTask) taskOptional.get()).getIOConfig().getMinimumMessageTime(), - ((KafkaIndexTask) taskOptional.get()).getIOConfig().getMaximumMessageTime() - ).equals(taskSequenceName); - } + return new KafkaRecordSupplier(spec.getIoConfig().getConsumerProperties(), sortingMapper); } - private ListenableFuture stopTasksInGroup(@Nullable TaskGroup taskGroup, String stopReasonFormat, Object... args) - { - if (taskGroup == null) { - return Futures.immediateFuture(null); - } - - log.info( - "Stopping all tasks in taskGroup[%s] because: [%s]", - taskGroup.groupId, - StringUtils.format(stopReasonFormat, args) + @Override + protected void scheduleReporting(ScheduledExecutorService reportingExec) + { + KafkaSupervisorIOConfig ioConfig = spec.getIoConfig(); + KafkaSupervisorTuningConfig tuningConfig = spec.getTuningConfig(); + reportingExec.scheduleAtFixedRate( + updateCurrentAndLatestOffsets(), + ioConfig.getStartDelay().getMillis() + INITIAL_GET_OFFSET_DELAY_MILLIS, // wait for tasks to start up + Math.max( + tuningConfig.getOffsetFetchPeriod().getMillis(), MINIMUM_GET_OFFSET_PERIOD_MILLIS + ), + TimeUnit.MILLISECONDS ); - final List> futures = new ArrayList<>(); - for (Entry entry : taskGroup.tasks.entrySet()) { - final String taskId = entry.getKey(); - final TaskData taskData = entry.getValue(); - if (taskData.status == null) { - killTask(taskId, "Killing task since task status is not known to supervisor"); - } else if (!taskData.status.isComplete()) { - futures.add(stopTask(taskId, false)); - } - } - - return Futures.successfulAsList(futures); - } - - private ListenableFuture stopTask(final String id, final boolean publish) - { - return Futures.transform( - taskClient.stopAsync(id, publish), new Function() - { - @Nullable - @Override - public Void apply(@Nullable Boolean result) - { - if (result == null || !result) { - killTask(id, "Task [%s] failed to stop in a timely manner, killing task", id); - } - return null; - } - } + reportingExec.scheduleAtFixedRate( + emitLag(), + ioConfig.getStartDelay().getMillis() + INITIAL_EMIT_LAG_METRIC_DELAY_MILLIS, // wait for tasks to start up + monitorSchedulerConfig.getEmitterPeriod().getMillis(), + TimeUnit.MILLISECONDS ); } - private void killTask(final String id, String reasonFormat, Object... args) + + @Override + protected int getTaskGroupIdForPartition(Integer partition) { - Optional taskQueue = taskMaster.getTaskQueue(); - if (taskQueue.isPresent()) { - taskQueue.get().shutdown(id, reasonFormat, args); - } else { - log.error("Failed to get task queue because I'm not the leader!"); - } + return partition % spec.getIoConfig().getTaskCount(); } - protected int getTaskGroupIdForPartition(int partition) + @Override + protected boolean checkSourceMetadataMatch(DataSourceMetadata metadata) { - return partition % ioConfig.getTaskCount(); + return metadata instanceof KafkaDataSourceMetadata; } - private boolean isTaskInPendingCompletionGroups(String taskId) + @Override + protected boolean doesTaskTypeMatchSupervisor(Task task) { - for (List taskGroups : pendingCompletionTaskGroups.values()) { - for (TaskGroup taskGroup : taskGroups) { - if (taskGroup.tasks.containsKey(taskId)) { - return true; - } - } - } - return false; + return task instanceof KafkaIndexTask; } - private SupervisorReport generateReport(boolean includeOffsets) + @Override + protected SeekableStreamSupervisorReportPayload createReportPayload( + int numPartitions, + boolean includeOffsets + ) { - int numPartitions = partitionGroups.values().stream().mapToInt(Map::size).sum(); - + KafkaSupervisorIOConfig ioConfig = spec.getIoConfig(); Map partitionLag = getLagPerPartition(getHighestCurrentOffsets()); - final KafkaSupervisorReportPayload payload = new KafkaSupervisorReportPayload( - dataSource, + return new KafkaSupervisorReportPayload( + spec.getDataSchema().getDataSource(), ioConfig.getTopic(), numPartitions, ioConfig.getReplicas(), ioConfig.getTaskDuration().getMillis() / 1000, - includeOffsets ? latestOffsetsFromKafka : null, + includeOffsets ? latestSequenceFromStream : null, includeOffsets ? partitionLag : null, includeOffsets ? partitionLag.values().stream().mapToLong(x -> Math.max(x, 0)).sum() : null, - includeOffsets ? offsetsLastUpdated : null, + includeOffsets ? sequenceLastUpdated : null, spec.isSuspended() ); - SupervisorReport report = new SupervisorReport<>( - dataSource, - DateTimes.nowUtc(), - payload - ); - - List taskReports = new ArrayList<>(); - - try { - for (TaskGroup taskGroup : taskGroups.values()) { - for (Entry entry : taskGroup.tasks.entrySet()) { - String taskId = entry.getKey(); - @Nullable - DateTime startTime = entry.getValue().startTime; - Map currentOffsets = entry.getValue().currentOffsets; - Long remainingSeconds = null; - if (startTime != null) { - long elapsedMillis = System.currentTimeMillis() - startTime.getMillis(); - long remainingMillis = Math.max(0, ioConfig.getTaskDuration().getMillis() - elapsedMillis); - remainingSeconds = TimeUnit.MILLISECONDS.toSeconds(remainingMillis); - } - - taskReports.add( - new TaskReportData( - taskId, - includeOffsets ? taskGroup.partitionOffsets : null, - includeOffsets ? currentOffsets : null, - startTime, - remainingSeconds, - TaskReportData.TaskType.ACTIVE, - includeOffsets ? getLagPerPartition(currentOffsets) : null - ) - ); - } - } - - for (List taskGroups : pendingCompletionTaskGroups.values()) { - for (TaskGroup taskGroup : taskGroups) { - for (Entry entry : taskGroup.tasks.entrySet()) { - String taskId = entry.getKey(); - @Nullable - DateTime startTime = entry.getValue().startTime; - Map currentOffsets = entry.getValue().currentOffsets; - Long remainingSeconds = null; - if (taskGroup.completionTimeout != null) { - remainingSeconds = Math.max(0, taskGroup.completionTimeout.getMillis() - System.currentTimeMillis()) - / 1000; - } - - taskReports.add( - new TaskReportData( - taskId, - includeOffsets ? taskGroup.partitionOffsets : null, - includeOffsets ? currentOffsets : null, - startTime, - remainingSeconds, - TaskReportData.TaskType.PUBLISHING, - null - ) - ); - } - } - } - - taskReports.forEach(payload::addTask); - } - catch (Exception e) { - log.warn(e, "Failed to generate status report"); - } - - return report; } - private Runnable buildRunTask() + + @Override + protected SeekableStreamIndexTaskIOConfig createTaskIoConfig( + int groupId, + Map startPartitions, + Map endPartitions, + String baseSequenceName, + DateTime minimumMessageTime, + DateTime maximumMessageTime, + Set exclusiveStartSequenceNumberPartitions, + SeekableStreamSupervisorIOConfig ioConfig + ) { - return () -> notices.add(new RunNotice()); + KafkaSupervisorIOConfig kafkaIoConfig = (KafkaSupervisorIOConfig) ioConfig; + return new KafkaIndexTaskIOConfig( + groupId, + baseSequenceName, + new SeekableStreamPartitions<>(kafkaIoConfig.getTopic(), startPartitions), + new SeekableStreamPartitions<>(kafkaIoConfig.getTopic(), endPartitions), + kafkaIoConfig.getConsumerProperties(), + true, + minimumMessageTime, + maximumMessageTime, + kafkaIoConfig.isSkipOffsetGaps() + ); } - private void updateLatestOffsetsFromKafka() + @Override + protected List> createIndexTasks( + int replicas, + String baseSequenceName, + ObjectMapper sortingMapper, + TreeMap> sequenceOffsets, + SeekableStreamIndexTaskIOConfig taskIoConfig, + SeekableStreamIndexTaskTuningConfig taskTuningConfig, + RowIngestionMetersFactory rowIngestionMetersFactory + ) throws JsonProcessingException { - synchronized (consumerLock) { - final List partitionInfoList = consumer.partitionsFor(ioConfig.getTopic()); - - if (partitionInfoList == null || partitionInfoList.size() == 0) { - throw new ISE("Could not retrieve partitions for topic [%s]", ioConfig.getTopic()); - } - - final Set topicPartitions = partitionInfoList - .stream() - .map(x -> new TopicPartition(x.topic(), x.partition())) - .collect(Collectors.toSet()); - - consumer.assign(topicPartitions); - - final Map endOffsets = consumer.endOffsets(topicPartitions); + final String checkpoints = sortingMapper.writerWithType(new TypeReference>>() + { + }).writeValueAsString(sequenceOffsets); + final Map context = spec.getContext() == null + ? ImmutableMap.of( + "checkpoints", + checkpoints, + IS_INCREMENTAL_HANDOFF_SUPPORTED, + true + ) : ImmutableMap.builder() + .put("checkpoints", checkpoints) + .put(IS_INCREMENTAL_HANDOFF_SUPPORTED, true) + .putAll(spec.getContext()) + .build(); - latestOffsetsFromKafka = endOffsets - .entrySet() - .stream() - .collect(Collectors.toMap(e -> e.getKey().partition(), Map.Entry::getValue)); + List> taskList = new ArrayList<>(); + for (int i = 0; i < replicas; i++) { + String taskId = Joiner.on("_").join(baseSequenceName, RandomIdUtils.getRandomId()); + taskList.add(new KafkaIndexTask( + taskId, + new TaskResource(baseSequenceName, 1), + spec.getDataSchema(), + (KafkaIndexTaskTuningConfig) taskTuningConfig, + (KafkaIndexTaskIOConfig) taskIoConfig, + context, + null, + null, + rowIngestionMetersFactory, + sortingMapper + )); } + return taskList; } - private Map getHighestCurrentOffsets() - { - return taskGroups - .values() - .stream() - .flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()) - .flatMap(taskData -> taskData.getValue().currentOffsets.entrySet().stream()) - .collect(Collectors.toMap(Entry::getKey, Entry::getValue, Long::max)); - } - private Map getLagPerPartition(Map currentOffsets) + @Override + protected Map getLagPerPartition(Map currentOffsets) { return currentOffsets .entrySet() @@ -2348,29 +271,42 @@ private Map getLagPerPartition(Map currentOffsets) .collect( Collectors.toMap( Entry::getKey, - e -> latestOffsetsFromKafka != null - && latestOffsetsFromKafka.get(e.getKey()) != null + e -> latestSequenceFromStream != null + && latestSequenceFromStream.get(e.getKey()) != null && e.getValue() != null - ? latestOffsetsFromKafka.get(e.getKey()) - e.getValue() + ? latestSequenceFromStream.get(e.getKey()) - e.getValue() : Integer.MIN_VALUE ) ); } + @Override + protected KafkaDataSourceMetadata createDataSourceMetaData(String topic, Map map) + { + return new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, map)); + } + + @Override + protected OrderedSequenceNumber makeSequenceNumber(Long seq, boolean isExclusive) + { + return KafkaSequenceNumber.of(seq); + } + private Runnable emitLag() { return () -> { try { Map highestCurrentOffsets = getHighestCurrentOffsets(); + String dataSource = spec.getDataSchema().getDataSource(); - if (latestOffsetsFromKafka == null) { + if (latestSequenceFromStream == null) { throw new ISE("Latest offsets from Kafka have not been fetched"); } - if (!latestOffsetsFromKafka.keySet().equals(highestCurrentOffsets.keySet())) { + if (!latestSequenceFromStream.keySet().equals(highestCurrentOffsets.keySet())) { log.warn( "Lag metric: Kafka partitions %s do not match task partitions %s", - latestOffsetsFromKafka.keySet(), + latestSequenceFromStream.keySet(), highestCurrentOffsets.keySet() ); } @@ -2401,239 +337,47 @@ private Runnable emitLag() }; } - private void updateCurrentOffsets() throws InterruptedException, ExecutionException, TimeoutException - { - final List> futures = Stream.concat( - taskGroups.values().stream().flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()), - pendingCompletionTaskGroups.values() - .stream() - .flatMap(List::stream) - .flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()) - ).map( - task -> Futures.transform( - taskClient.getCurrentOffsetsAsync(task.getKey(), false), - (Function, Void>) (currentOffsets) -> { - - if (currentOffsets != null && !currentOffsets.isEmpty()) { - task.getValue().currentOffsets = currentOffsets; - } - - return null; - } - ) - ).collect(Collectors.toList()); - - Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - } - - @VisibleForTesting - Runnable updateCurrentAndLatestOffsets() + @Override + protected Long getNotSetMarker() { - return () -> { - try { - updateCurrentOffsets(); - updateLatestOffsetsFromKafka(); - offsetsLastUpdated = DateTimes.nowUtc(); - } - catch (Exception e) { - log.warn(e, "Exception while getting current/latest offsets"); - } - }; + return NOT_SET; } - /** - * Collect row ingestion stats from all tasks managed by this supervisor. - * - * @return A map of groupId->taskId->task row stats - * - * @throws InterruptedException - * @throws ExecutionException - * @throws TimeoutException - */ - private Map> getCurrentTotalStats() - throws InterruptedException, ExecutionException, TimeoutException + @Override + protected Long getEndOfPartitionMarker() { - Map> allStats = new HashMap<>(); - final List> futures = new ArrayList<>(); - final List> groupAndTaskIds = new ArrayList<>(); - - for (int groupId : taskGroups.keySet()) { - TaskGroup group = taskGroups.get(groupId); - for (String taskId : group.taskIds()) { - futures.add( - Futures.transform( - taskClient.getMovingAveragesAsync(taskId), - (Function, StatsFromTaskResult>) (currentStats) -> { - return new StatsFromTaskResult( - groupId, - taskId, - currentStats - ); - } - ) - ); - groupAndTaskIds.add(new Pair<>(groupId, taskId)); - } - } - - for (int groupId : pendingCompletionTaskGroups.keySet()) { - List pendingGroups = pendingCompletionTaskGroups.get(groupId); - for (TaskGroup pendingGroup : pendingGroups) { - for (String taskId : pendingGroup.taskIds()) { - futures.add( - Futures.transform( - taskClient.getMovingAveragesAsync(taskId), - (Function, StatsFromTaskResult>) (currentStats) -> new StatsFromTaskResult( - groupId, - taskId, - currentStats - ) - ) - ); - groupAndTaskIds.add(new Pair<>(groupId, taskId)); - } - } - } - - List results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); - for (int i = 0; i < results.size(); i++) { - StatsFromTaskResult result = results.get(i); - if (result != null) { - Map groupMap = allStats.computeIfAbsent(result.getGroupId(), k -> new HashMap<>()); - groupMap.put(result.getTaskId(), result.getStats()); - } else { - Pair groupAndTaskId = groupAndTaskIds.get(i); - log.error("Failed to get stats for group[%d]-task[%s]", groupAndTaskId.lhs, groupAndTaskId.rhs); - } - } - - return allStats; + return END_OF_PARTITION; } - @VisibleForTesting - void addTaskGroupToActivelyReadingTaskGroup( - int taskGroupId, - ImmutableMap partitionOffsets, - Optional minMsgTime, - Optional maxMsgTime, - Set tasks - ) + @Override + protected boolean isEndOfShard(Long seqNum) { - TaskGroup group = new TaskGroup( - taskGroupId, - partitionOffsets, - minMsgTime, - maxMsgTime - ); - group.tasks.putAll(tasks.stream().collect(Collectors.toMap(x -> x, x -> new TaskData()))); - if (taskGroups.putIfAbsent(taskGroupId, group) != null) { - throw new ISE( - "trying to add taskGroup with ID [%s] to actively reading task groups, but group already exists.", - taskGroupId - ); - } + return false; } - @VisibleForTesting - void addTaskGroupToPendingCompletionTaskGroup( - int taskGroupId, - ImmutableMap partitionOffsets, - Optional minMsgTime, - Optional maxMsgTime, - Set tasks + @Override + protected void updateLatestSequenceFromStream( + RecordSupplier recordSupplier, + Set> partitions ) { - TaskGroup group = new TaskGroup( - taskGroupId, - partitionOffsets, - minMsgTime, - maxMsgTime - ); - group.tasks.putAll(tasks.stream().collect(Collectors.toMap(x -> x, x -> new TaskData()))); - pendingCompletionTaskGroups.computeIfAbsent(taskGroupId, x -> new CopyOnWriteArrayList<>()) - .add(group); - } - - @VisibleForTesting - @Nullable - TaskGroup removeTaskGroup(int taskGroupId) - { - return taskGroups.remove(taskGroupId); - } - - @VisibleForTesting - void moveTaskGroupToPendingCompletion(int taskGroupId) - { - final TaskGroup taskGroup = taskGroups.remove(taskGroupId); - if (taskGroup != null) { - pendingCompletionTaskGroups.computeIfAbsent(taskGroupId, k -> new CopyOnWriteArrayList<>()).add(taskGroup); - } - } - - @VisibleForTesting - int getNoticesQueueSize() - { - return notices.size(); - } - - private static class StatsFromTaskResult - { - private final String groupId; - private final String taskId; - private final Map stats; - - public StatsFromTaskResult( - int groupId, - String taskId, - Map stats - ) - { - this.groupId = String.valueOf(groupId); - this.taskId = taskId; - this.stats = stats; - } - - public String getGroupId() - { - return groupId; - } - - public String getTaskId() - { - return taskId; - } - - public Map getStats() - { - return stats; - } - } - - // exposed for testing for visibility into initialization state - @VisibleForTesting - public boolean isStarted() - { - return started; - } - - // exposed for testing for visibility into initialization state - @VisibleForTesting - public boolean isLifecycleStarted() - { - return lifecycleStarted; + latestSequenceFromStream = partitions.stream() + .collect(Collectors.toMap( + StreamPartition::getPartitionId, + recordSupplier::getPosition + )); } - // exposed for testing for visibility into initialization state - @VisibleForTesting - public int getInitRetryCounter() + @Override + protected String baseTaskName() { - return initRetryCounter; + return "index_kafka"; } - // exposed for testing to allow "bootstrap.servers" to be changed after supervisor is created + @Override @VisibleForTesting public KafkaSupervisorIOConfig getIoConfig() { - return ioConfig; + return spec.getIoConfig(); } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java index 44c2bb2d6f73..80b842b83496 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java @@ -21,32 +21,21 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.common.base.Optional; import com.google.common.base.Preconditions; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig; import org.apache.druid.java.util.common.StringUtils; -import org.joda.time.Duration; import org.joda.time.Period; import java.util.Map; -public class KafkaSupervisorIOConfig +public class KafkaSupervisorIOConfig extends SeekableStreamSupervisorIOConfig { public static final String BOOTSTRAP_SERVERS_KEY = "bootstrap.servers"; public static final String TRUST_STORE_PASSWORD_KEY = "ssl.truststore.password"; public static final String KEY_STORE_PASSWORD_KEY = "ssl.keystore.password"; public static final String KEY_PASSWORD_KEY = "ssl.key.password"; - private final String topic; - private final Integer replicas; - private final Integer taskCount; - private final Duration taskDuration; private final Map consumerProperties; - private final Duration startDelay; - private final Duration period; - private final boolean useEarliestOffset; - private final Duration completionTimeout; - private final Optional lateMessageRejectionPeriod; - private final Optional earlyMessageRejectionPeriod; private final boolean skipOffsetGaps; @JsonCreator @@ -65,51 +54,31 @@ public KafkaSupervisorIOConfig( @JsonProperty("skipOffsetGaps") Boolean skipOffsetGaps ) { - this.topic = Preconditions.checkNotNull(topic, "topic"); + super( + Preconditions.checkNotNull(topic, "topic"), + replicas, + taskCount, + taskDuration, + startDelay, + period, + useEarliestOffset, + completionTimeout, + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod + ); + this.consumerProperties = Preconditions.checkNotNull(consumerProperties, "consumerProperties"); Preconditions.checkNotNull( consumerProperties.get(BOOTSTRAP_SERVERS_KEY), StringUtils.format("consumerProperties must contain entry for [%s]", BOOTSTRAP_SERVERS_KEY) ); - - this.replicas = replicas != null ? replicas : 1; - this.taskCount = taskCount != null ? taskCount : 1; - this.taskDuration = defaultDuration(taskDuration, "PT1H"); - this.startDelay = defaultDuration(startDelay, "PT5S"); - this.period = defaultDuration(period, "PT30S"); - this.useEarliestOffset = useEarliestOffset != null ? useEarliestOffset : false; - this.completionTimeout = defaultDuration(completionTimeout, "PT30M"); - this.lateMessageRejectionPeriod = lateMessageRejectionPeriod == null - ? Optional.absent() - : Optional.of(lateMessageRejectionPeriod.toStandardDuration()); - this.earlyMessageRejectionPeriod = earlyMessageRejectionPeriod == null - ? Optional.absent() - : Optional.of(earlyMessageRejectionPeriod.toStandardDuration()); this.skipOffsetGaps = skipOffsetGaps != null ? skipOffsetGaps : false; } @JsonProperty public String getTopic() { - return topic; - } - - @JsonProperty - public Integer getReplicas() - { - return replicas; - } - - @JsonProperty - public Integer getTaskCount() - { - return taskCount; - } - - @JsonProperty - public Duration getTaskDuration() - { - return taskDuration; + return getStream(); } @JsonProperty @@ -118,40 +87,10 @@ public Map getConsumerProperties() return consumerProperties; } - @JsonProperty - public Duration getStartDelay() - { - return startDelay; - } - - @JsonProperty - public Duration getPeriod() - { - return period; - } - @JsonProperty public boolean isUseEarliestOffset() { - return useEarliestOffset; - } - - @JsonProperty - public Duration getCompletionTimeout() - { - return completionTimeout; - } - - @JsonProperty - public Optional getEarlyMessageRejectionPeriod() - { - return earlyMessageRejectionPeriod; - } - - @JsonProperty - public Optional getLateMessageRejectionPeriod() - { - return lateMessageRejectionPeriod; + return isUseEarliestSequenceNumber(); } @JsonProperty @@ -164,22 +103,19 @@ public boolean isSkipOffsetGaps() public String toString() { return "KafkaSupervisorIOConfig{" + - "topic='" + topic + '\'' + - ", replicas=" + replicas + - ", taskCount=" + taskCount + - ", taskDuration=" + taskDuration + + "topic='" + getTopic() + '\'' + + ", replicas=" + getReplicas() + + ", taskCount=" + getTaskCount() + + ", taskDuration=" + getTaskDuration() + ", consumerProperties=" + consumerProperties + - ", startDelay=" + startDelay + - ", period=" + period + - ", useEarliestOffset=" + useEarliestOffset + - ", completionTimeout=" + completionTimeout + - ", lateMessageRejectionPeriod=" + lateMessageRejectionPeriod + + ", startDelay=" + getStartDelay() + + ", period=" + getPeriod() + + ", useEarliestOffset=" + isUseEarliestOffset() + + ", completionTimeout=" + getCompletionTimeout() + + ", earlyMessageRejectionPeriod=" + getEarlyMessageRejectionPeriod() + + ", lateMessageRejectionPeriod=" + getLateMessageRejectionPeriod() + ", skipOffsetGaps=" + skipOffsetGaps + '}'; } - private static Duration defaultDuration(final Period period, final String theDefault) - { - return (period == null ? new Period(theDefault) : period).toStandardDuration(); - } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorReportPayload.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorReportPayload.java index d9533a37fb24..d5f4efa02a20 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorReportPayload.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorReportPayload.java @@ -19,30 +19,14 @@ package org.apache.druid.indexing.kafka.supervisor; -import com.fasterxml.jackson.annotation.JsonInclude; -import com.fasterxml.jackson.annotation.JsonProperty; -import org.apache.druid.java.util.common.IAE; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorReportPayload; import org.joda.time.DateTime; import javax.annotation.Nullable; -import java.util.ArrayList; -import java.util.List; import java.util.Map; -public class KafkaSupervisorReportPayload +public class KafkaSupervisorReportPayload extends SeekableStreamSupervisorReportPayload { - private final String dataSource; - private final String topic; - private final int partitions; - private final int replicas; - private final long durationSeconds; - private final List activeTasks; - private final List publishingTasks; - private final Map latestOffsets; - private final Map minimumLag; - private final Long aggregateLag; - private final DateTime offsetsLastUpdated; - private final boolean suspended; public KafkaSupervisorReportPayload( String dataSource, @@ -57,122 +41,38 @@ public KafkaSupervisorReportPayload( boolean suspended ) { - this.dataSource = dataSource; - this.topic = topic; - this.partitions = partitions; - this.replicas = replicas; - this.durationSeconds = durationSeconds; - this.activeTasks = new ArrayList<>(); - this.publishingTasks = new ArrayList<>(); - this.latestOffsets = latestOffsets; - this.minimumLag = minimumLag; - this.aggregateLag = aggregateLag; - this.offsetsLastUpdated = offsetsLastUpdated; - this.suspended = suspended; + super( + dataSource, + topic, + partitions, + replicas, + durationSeconds, + latestOffsets, + minimumLag, + aggregateLag, + offsetsLastUpdated, + suspended + ); } - public void addTask(TaskReportData data) - { - if (data.getType().equals(TaskReportData.TaskType.ACTIVE)) { - activeTasks.add(data); - } else if (data.getType().equals(TaskReportData.TaskType.PUBLISHING)) { - publishingTasks.add(data); - } else { - throw new IAE("Unknown task type [%s]", data.getType().name()); - } - } - - @JsonProperty - public String getDataSource() - { - return dataSource; - } - - @JsonProperty - public String getTopic() - { - return topic; - } - - @JsonProperty - public int getPartitions() - { - return partitions; - } - - @JsonProperty - public int getReplicas() - { - return replicas; - } - - @JsonProperty - public long getDurationSeconds() - { - return durationSeconds; - } - - @JsonProperty - public List getActiveTasks() - { - return activeTasks; - } - - @JsonProperty - public List getPublishingTasks() - { - return publishingTasks; - } - - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_NULL) - public Map getLatestOffsets() - { - return latestOffsets; - } - - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_NULL) - public Map getMinimumLag() - { - return minimumLag; - } - - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_NULL) - public Long getAggregateLag() - { - return aggregateLag; - } - - @JsonProperty - public DateTime getOffsetsLastUpdated() - { - return offsetsLastUpdated; - } - - @JsonProperty - public boolean getSuspended() - { - return suspended; - } @Override public String toString() { - return "{" + - "dataSource='" + dataSource + '\'' + - ", topic='" + topic + '\'' + - ", partitions=" + partitions + - ", replicas=" + replicas + - ", durationSeconds=" + durationSeconds + - ", active=" + activeTasks + - ", publishing=" + publishingTasks + - (latestOffsets != null ? ", latestOffsets=" + latestOffsets : "") + - (minimumLag != null ? ", minimumLag=" + minimumLag : "") + - (aggregateLag != null ? ", aggregateLag=" + aggregateLag : "") + - (offsetsLastUpdated != null ? ", offsetsLastUpdated=" + offsetsLastUpdated : "") + - ", suspended=" + suspended + + return "KafkaSupervisorReportPayload{" + + "dataSource='" + getDataSource() + '\'' + + ", topic='" + getStream() + '\'' + + ", partitions=" + getPartitions() + + ", replicas=" + getReplicas() + + ", durationSeconds=" + getDurationSeconds() + + ", active=" + getActiveTasks() + + ", publishing=" + getPublishingTasks() + + (getLatestOffsets() != null ? ", latestOffsets=" + getLatestOffsets() : "") + + (getMinimumLag() != null ? ", minimumLag=" + getMinimumLag() : "") + + (getAggregateLag() != null ? ", aggregateLag=" + getAggregateLag() : "") + + (getOffsetsLastUpdated() != null ? ", sequenceLastUpdated=" + getOffsetsLastUpdated() : "") + + ", suspended=" + getSuspended() + '}'; } + } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java index 28dab02c56e2..33ad3e11bee7 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java @@ -23,8 +23,6 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; import org.apache.druid.guice.annotations.Json; import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; import org.apache.druid.indexing.kafka.KafkaIndexTaskClientFactory; @@ -32,30 +30,15 @@ import org.apache.druid.indexing.overlord.TaskMaster; import org.apache.druid.indexing.overlord.TaskStorage; import org.apache.druid.indexing.overlord.supervisor.Supervisor; -import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorSpec; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.segment.indexing.DataSchema; import org.apache.druid.server.metrics.DruidMonitorSchedulerConfig; -import java.util.List; import java.util.Map; -public class KafkaSupervisorSpec implements SupervisorSpec +public class KafkaSupervisorSpec extends SeekableStreamSupervisorSpec { - private final DataSchema dataSchema; - private final KafkaSupervisorTuningConfig tuningConfig; - private final KafkaSupervisorIOConfig ioConfig; - private final Map context; - - private final TaskStorage taskStorage; - private final TaskMaster taskMaster; - private final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator; - private final KafkaIndexTaskClientFactory kafkaIndexTaskClientFactory; - private final ObjectMapper mapper; - private final ServiceEmitter emitter; - private final DruidMonitorSchedulerConfig monitorSchedulerConfig; - private final RowIngestionMetersFactory rowIngestionMetersFactory; - private final boolean suspended; @JsonCreator public KafkaSupervisorSpec( @@ -74,93 +57,47 @@ public KafkaSupervisorSpec( @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory ) { - this.dataSchema = Preconditions.checkNotNull(dataSchema, "dataSchema"); - this.tuningConfig = tuningConfig != null - ? tuningConfig - : new KafkaSupervisorTuningConfig( - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null, - null - ); - this.ioConfig = Preconditions.checkNotNull(ioConfig, "ioConfig"); - this.context = context; - - this.taskStorage = taskStorage; - this.taskMaster = taskMaster; - this.indexerMetadataStorageCoordinator = indexerMetadataStorageCoordinator; - this.kafkaIndexTaskClientFactory = kafkaIndexTaskClientFactory; - this.mapper = mapper; - this.emitter = emitter; - this.monitorSchedulerConfig = monitorSchedulerConfig; - this.rowIngestionMetersFactory = rowIngestionMetersFactory; - this.suspended = suspended != null ? suspended : false; - } - - @JsonProperty - public DataSchema getDataSchema() - { - return dataSchema; - } - - @JsonProperty - public KafkaSupervisorTuningConfig getTuningConfig() - { - return tuningConfig; - } - - @JsonProperty - public KafkaSupervisorIOConfig getIoConfig() - { - return ioConfig; - } - - @JsonProperty - public Map getContext() - { - return context; - } - - @Override - @JsonProperty("suspended") - public boolean isSuspended() - { - return suspended; - } - - public ServiceEmitter getEmitter() - { - return emitter; - } - - @Override - public String getId() - { - return dataSchema.getDataSource(); - } - - public DruidMonitorSchedulerConfig getMonitorSchedulerConfig() - { - return monitorSchedulerConfig; + super( + dataSchema, + tuningConfig != null + ? tuningConfig + : new KafkaSupervisorTuningConfig( + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ), + ioConfig, + context, + suspended, + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + kafkaIndexTaskClientFactory, + mapper, + emitter, + monitorSchedulerConfig, + rowIngestionMetersFactory + ); } @Override @@ -170,7 +107,7 @@ public Supervisor createSupervisor() taskStorage, taskMaster, indexerMetadataStorageCoordinator, - kafkaIndexTaskClientFactory, + (KafkaIndexTaskClientFactory) indexTaskClientFactory, mapper, this, rowIngestionMetersFactory @@ -178,49 +115,48 @@ public Supervisor createSupervisor() } @Override - public List getDataSources() - { - return ImmutableList.of(getDataSchema().getDataSource()); - } - - @Override - public String toString() + @JsonProperty + public KafkaSupervisorTuningConfig getTuningConfig() { - return "KafkaSupervisorSpec{" + - "dataSchema=" + dataSchema + - ", tuningConfig=" + tuningConfig + - ", ioConfig=" + ioConfig + - '}'; + return (KafkaSupervisorTuningConfig) super.getTuningConfig(); } @Override - public KafkaSupervisorSpec createSuspendedSpec() + @JsonProperty + public KafkaSupervisorIOConfig getIoConfig() { - return toggleSuspend(true); + return (KafkaSupervisorIOConfig) super.getIoConfig(); } @Override - public KafkaSupervisorSpec createRunningSpec() - { - return toggleSuspend(false); - } - - private KafkaSupervisorSpec toggleSuspend(boolean suspend) + protected KafkaSupervisorSpec toggleSuspend(boolean suspend) { return new KafkaSupervisorSpec( - dataSchema, - tuningConfig, - ioConfig, - context, + getDataSchema(), + getTuningConfig(), + getIoConfig(), + getContext(), suspend, taskStorage, taskMaster, indexerMetadataStorageCoordinator, - kafkaIndexTaskClientFactory, + (KafkaIndexTaskClientFactory) indexTaskClientFactory, mapper, emitter, monitorSchedulerConfig, rowIngestionMetersFactory ); } + + @Override + public String toString() + { + return "KafkaSupervisorSpec{" + + "dataSchema=" + getDataSchema() + + ", tuningConfig=" + getTuningConfig() + + ", ioConfig=" + getIoConfig() + + ", context=" + getContext() + + ", suspend=" + isSuspended() + + '}'; + } } diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfig.java b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfig.java index 4715f336f759..8e4c6e9f9137 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfig.java +++ b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfig.java @@ -20,7 +20,9 @@ package org.apache.druid.indexing.kafka.supervisor; import com.fasterxml.jackson.annotation.JsonProperty; -import org.apache.druid.indexing.kafka.KafkaTuningConfig; +import org.apache.druid.indexing.kafka.KafkaIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorTuningConfig; import org.apache.druid.segment.IndexSpec; import org.apache.druid.segment.indexing.TuningConfigs; import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory; @@ -30,8 +32,11 @@ import javax.annotation.Nullable; import java.io.File; -public class KafkaSupervisorTuningConfig extends KafkaTuningConfig +public class KafkaSupervisorTuningConfig extends KafkaIndexTaskTuningConfig + implements SeekableStreamSupervisorTuningConfig { + private static final String DEFAULT_OFFSET_FETCH_PERIOD = "PT30S"; + private final Integer workerThreads; private final Integer chatThreads; private final Long chatRetries; @@ -85,39 +90,49 @@ public KafkaSupervisorTuningConfig( maxParseExceptions, maxSavedParseExceptions ); - this.workerThreads = workerThreads; this.chatThreads = chatThreads; - this.chatRetries = (chatRetries != null ? chatRetries : 8); - this.httpTimeout = defaultDuration(httpTimeout, "PT10S"); - this.shutdownTimeout = defaultDuration(shutdownTimeout, "PT80S"); - this.offsetFetchPeriod = defaultDuration(offsetFetchPeriod, "PT30S"); + this.chatRetries = (chatRetries != null ? chatRetries : DEFAULT_CHAT_RETRIES); + this.httpTimeout = SeekableStreamSupervisorTuningConfig.defaultDuration(httpTimeout, DEFAULT_HTTP_TIMEOUT); + this.shutdownTimeout = SeekableStreamSupervisorTuningConfig.defaultDuration( + shutdownTimeout, + DEFAULT_SHUTDOWN_TIMEOUT + ); + this.offsetFetchPeriod = SeekableStreamSupervisorTuningConfig.defaultDuration( + offsetFetchPeriod, + DEFAULT_OFFSET_FETCH_PERIOD + ); } + @Override @JsonProperty public Integer getWorkerThreads() { return workerThreads; } + @Override @JsonProperty public Integer getChatThreads() { return chatThreads; } + @Override @JsonProperty public Long getChatRetries() { return chatRetries; } + @Override @JsonProperty public Duration getHttpTimeout() { return httpTimeout; } + @Override @JsonProperty public Duration getShutdownTimeout() { @@ -159,8 +174,27 @@ public String toString() '}'; } - private static Duration defaultDuration(final Period period, final String theDefault) + @Override + public SeekableStreamIndexTaskTuningConfig convertToTaskTuningConfig() { - return (period == null ? new Period(theDefault) : period).toStandardDuration(); + return new KafkaIndexTaskTuningConfig( + getMaxRowsInMemory(), + getMaxBytesInMemory(), + getMaxRowsPerSegment(), + getMaxTotalRows(), + getIntermediatePersistPeriod(), + getBasePersistDirectory(), + getMaxPendingPersists(), + getIndexSpec(), + true, + isReportParseExceptions(), + getHandoffConditionTimeout(), + isResetOffsetAutomatically(), + getSegmentWriteOutMediumFactory(), + getIntermediateHandoffPeriod(), + isLogParseExceptions(), + getMaxParseExceptions(), + getMaxSavedParseExceptions() + ); } } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadataTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadataTest.java index 51afdeb8ea40..89f5ce8b0e76 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadataTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaDataSourceMetadataTest.java @@ -20,6 +20,7 @@ package org.apache.druid.indexing.kafka; import com.google.common.collect.ImmutableMap; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; import org.junit.Assert; import org.junit.Test; @@ -125,6 +126,6 @@ public void testMinus() private static KafkaDataSourceMetadata KM(String topic, Map offsets) { - return new KafkaDataSourceMetadata(new KafkaPartitions(topic, offsets)); + return new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, offsets)); } } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIOConfigTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIOConfigTest.java index 2321974ff383..7ce8df024443 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIOConfigTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIOConfigTest.java @@ -32,6 +32,8 @@ import org.junit.Test; import org.junit.rules.ExpectedException; +import java.util.Collections; + public class KafkaIOConfigTest { private final ObjectMapper mapper; @@ -57,7 +59,7 @@ public void testSerdeWithDefaults() throws Exception + " \"consumerProperties\": {\"bootstrap.servers\":\"localhost:9092\"}\n" + "}"; - KafkaIOConfig config = (KafkaIOConfig) mapper.readValue( + KafkaIndexTaskIOConfig config = (KafkaIndexTaskIOConfig) mapper.readValue( mapper.writeValueAsString( mapper.readValue( jsonStr, @@ -67,15 +69,16 @@ public void testSerdeWithDefaults() throws Exception ); Assert.assertEquals("my-sequence-name", config.getBaseSequenceName()); - Assert.assertEquals("mytopic", config.getStartPartitions().getTopic()); - Assert.assertEquals(ImmutableMap.of(0, 1L, 1, 10L), config.getStartPartitions().getPartitionOffsetMap()); - Assert.assertEquals("mytopic", config.getEndPartitions().getTopic()); - Assert.assertEquals(ImmutableMap.of(0, 15L, 1, 200L), config.getEndPartitions().getPartitionOffsetMap()); + Assert.assertEquals("mytopic", config.getStartPartitions().getStream()); + Assert.assertEquals(ImmutableMap.of(0, 1L, 1, 10L), config.getStartPartitions().getPartitionSequenceNumberMap()); + Assert.assertEquals("mytopic", config.getEndPartitions().getStream()); + Assert.assertEquals(ImmutableMap.of(0, 15L, 1, 200L), config.getEndPartitions().getPartitionSequenceNumberMap()); Assert.assertEquals(ImmutableMap.of("bootstrap.servers", "localhost:9092"), config.getConsumerProperties()); Assert.assertTrue(config.isUseTransaction()); Assert.assertFalse("minimumMessageTime", config.getMinimumMessageTime().isPresent()); Assert.assertFalse("maximumMessageTime", config.getMaximumMessageTime().isPresent()); Assert.assertFalse("skipOffsetGaps", config.isSkipOffsetGaps()); + Assert.assertEquals(Collections.EMPTY_SET, config.getExclusiveStartSequenceNumberPartitions()); } @Test @@ -94,7 +97,7 @@ public void testSerdeWithNonDefaults() throws Exception + " \"skipOffsetGaps\": true\n" + "}"; - KafkaIOConfig config = (KafkaIOConfig) mapper.readValue( + KafkaIndexTaskIOConfig config = (KafkaIndexTaskIOConfig) mapper.readValue( mapper.writeValueAsString( mapper.readValue( jsonStr, @@ -104,15 +107,17 @@ public void testSerdeWithNonDefaults() throws Exception ); Assert.assertEquals("my-sequence-name", config.getBaseSequenceName()); - Assert.assertEquals("mytopic", config.getStartPartitions().getTopic()); - Assert.assertEquals(ImmutableMap.of(0, 1L, 1, 10L), config.getStartPartitions().getPartitionOffsetMap()); - Assert.assertEquals("mytopic", config.getEndPartitions().getTopic()); - Assert.assertEquals(ImmutableMap.of(0, 15L, 1, 200L), config.getEndPartitions().getPartitionOffsetMap()); + Assert.assertEquals("mytopic", config.getStartPartitions().getStream()); + Assert.assertEquals(ImmutableMap.of(0, 1L, 1, 10L), config.getStartPartitions().getPartitionSequenceNumberMap()); + Assert.assertEquals("mytopic", config.getEndPartitions().getStream()); + Assert.assertEquals(ImmutableMap.of(0, 15L, 1, 200L), config.getEndPartitions().getPartitionSequenceNumberMap()); Assert.assertEquals(ImmutableMap.of("bootstrap.servers", "localhost:9092"), config.getConsumerProperties()); Assert.assertFalse(config.isUseTransaction()); Assert.assertEquals(DateTimes.of("2016-05-31T12:00Z"), config.getMinimumMessageTime().get()); Assert.assertEquals(DateTimes.of("2016-05-31T14:00Z"), config.getMaximumMessageTime().get()); Assert.assertTrue("skipOffsetGaps", config.isSkipOffsetGaps()); + Assert.assertEquals(Collections.EMPTY_SET, config.getExclusiveStartSequenceNumberPartitions()); + } @Test @@ -212,7 +217,7 @@ public void testStartAndEndTopicMatch() throws Exception exception.expect(JsonMappingException.class); exception.expectCause(CoreMatchers.isA(IllegalArgumentException.class)); - exception.expectMessage(CoreMatchers.containsString("start topic and end topic must match")); + exception.expectMessage(CoreMatchers.containsString("start topic/stream and end topic/stream must match")); mapper.readValue(jsonStr, IOConfig.class); } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientTest.java index 5f0d915e5c8d..28cc0bbe9d15 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientTest.java @@ -23,7 +23,6 @@ import com.google.common.base.Optional; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; @@ -31,6 +30,7 @@ import org.apache.druid.indexer.TaskStatus; import org.apache.druid.indexing.common.IndexTaskClient; import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.IAE; @@ -41,6 +41,7 @@ import org.apache.druid.java.util.http.client.response.FullResponseHolder; import org.easymock.Capture; import org.easymock.CaptureType; +import org.easymock.EasyMock; import org.easymock.EasyMockSupport; import org.jboss.netty.handler.codec.http.HttpHeaders; import org.jboss.netty.handler.codec.http.HttpMethod; @@ -60,15 +61,12 @@ import java.io.IOException; import java.net.URL; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; -import static org.easymock.EasyMock.anyObject; -import static org.easymock.EasyMock.capture; -import static org.easymock.EasyMock.eq; import static org.easymock.EasyMock.expect; -import static org.easymock.EasyMock.reset; @RunWith(Parameterized.class) public class KafkaIndexTaskClientTest extends EasyMockSupport @@ -78,7 +76,7 @@ public class KafkaIndexTaskClientTest extends EasyMockSupport private static final ObjectMapper objectMapper = new DefaultObjectMapper(); private static final String TEST_ID = "test-id"; - private static final List TEST_IDS = Lists.newArrayList("test-id1", "test-id2", "test-id3", "test-id4"); + private static final List TEST_IDS = Arrays.asList("test-id1", "test-id2", "test-id3", "test-id4"); private static final String TEST_HOST = "test-host"; private static final int TEST_PORT = 1234; private static final int TEST_TLS_PORT = -1; @@ -116,9 +114,12 @@ public void setUp() headers = createMock(HttpHeaders.class); client = new TestableKafkaIndexTaskClient(httpClient, objectMapper, taskInfoProvider); - expect(taskInfoProvider.getTaskLocation(TEST_ID)).andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) - .anyTimes(); - expect(taskInfoProvider.getTaskStatus(TEST_ID)).andReturn(Optional.of(TaskStatus.running(TEST_ID))).anyTimes(); + expect(taskInfoProvider.getTaskLocation(TEST_ID)) + .andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) + .anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.running(TEST_ID))) + .anyTimes(); for (String testId : TEST_IDS) { expect(taskInfoProvider.getTaskLocation(testId)) @@ -139,16 +140,18 @@ public void tearDown() @Test public void testNoTaskLocation() throws IOException { - reset(taskInfoProvider); + EasyMock.reset(taskInfoProvider); expect(taskInfoProvider.getTaskLocation(TEST_ID)).andReturn(TaskLocation.unknown()).anyTimes(); - expect(taskInfoProvider.getTaskStatus(TEST_ID)).andReturn(Optional.of(TaskStatus.running(TEST_ID))).anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.running(TEST_ID))) + .anyTimes(); replayAll(); Assert.assertEquals(false, client.stop(TEST_ID, true)); Assert.assertEquals(false, client.resume(TEST_ID)); Assert.assertEquals(ImmutableMap.of(), client.pause(TEST_ID)); Assert.assertEquals(ImmutableMap.of(), client.pause(TEST_ID)); - Assert.assertEquals(KafkaIndexTask.Status.NOT_STARTED, client.getStatus(TEST_ID)); + Assert.assertEquals(Status.NOT_STARTED, client.getStatus(TEST_ID)); Assert.assertEquals(null, client.getStartTime(TEST_ID)); Assert.assertEquals(ImmutableMap.of(), client.getCurrentOffsets(TEST_ID, true)); Assert.assertEquals(ImmutableMap.of(), client.getEndOffsets(TEST_ID)); @@ -164,10 +167,13 @@ public void testTaskNotRunnableException() expectedException.expect(IndexTaskClient.TaskNotRunnableException.class); expectedException.expectMessage("Aborting request because task [test-id] is not runnable"); - reset(taskInfoProvider); - expect(taskInfoProvider.getTaskLocation(TEST_ID)).andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) - .anyTimes(); - expect(taskInfoProvider.getTaskStatus(TEST_ID)).andReturn(Optional.of(TaskStatus.failure(TEST_ID))).anyTimes(); + EasyMock.reset(taskInfoProvider); + expect(taskInfoProvider.getTaskLocation(TEST_ID)) + .andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) + .anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.failure(TEST_ID))) + .anyTimes(); replayAll(); client.getCurrentOffsets(TEST_ID, true); @@ -183,9 +189,9 @@ public void testInternalServerError() expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.INTERNAL_SERVER_ERROR).times(2); expect( httpClient.go( - anyObject(Request.class), - anyObject(FullResponseHandler.class), - eq(TEST_HTTP_TIMEOUT) + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) ) ).andReturn( Futures.immediateFuture(responseHolder) @@ -206,9 +212,9 @@ public void testBadRequest() expect(responseHolder.getContent()).andReturn(""); expect( httpClient.go( - anyObject(Request.class), - anyObject(FullResponseHandler.class), - eq(TEST_HTTP_TIMEOUT) + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) ) ).andReturn( Futures.immediateFuture(responseHolder) @@ -231,9 +237,9 @@ public void testTaskLocationMismatch() expect(headers.get("X-Druid-Task-Id")).andReturn("a-different-task-id"); expect( httpClient.go( - anyObject(Request.class), - anyObject(FullResponseHandler.class), - eq(TEST_HTTP_TIMEOUT) + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) ) ).andReturn( Futures.immediateFuture(responseHolder) @@ -252,7 +258,11 @@ public void testGetCurrentOffsets() throws Exception Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}"); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -287,7 +297,11 @@ public void testGetCurrentOffsetsWithRetry() throws Exception expect(response.headers()).andReturn(headers).times(2); expect(headers.get("X-Druid-Task-Id")).andReturn(TEST_ID).times(2); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(3); @@ -327,9 +341,9 @@ public void testGetCurrentOffsetsWithExhaustedRetries() expect( httpClient.go( - anyObject(Request.class), - anyObject(FullResponseHandler.class), - eq(TEST_HTTP_TIMEOUT) + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) ) ).andReturn(Futures.immediateFuture(responseHolder)).anyTimes(); replayAll(); @@ -344,7 +358,11 @@ public void testGetEndOffsets() throws Exception Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}"); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -378,7 +396,11 @@ public void testGetStartTime() throws Exception expect(response.headers()).andReturn(headers); expect(headers.get("X-Druid-Task-Id")).andReturn(null); expect(responseHolder.getContent()).andReturn(String.valueOf(now.getMillis())).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(2); replayAll(); @@ -400,17 +422,21 @@ public void testGetStartTime() throws Exception @Test public void testGetStatus() throws Exception { - KafkaIndexTask.Status status = KafkaIndexTask.Status.READING; + Status status = Status.READING; Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); expect(responseHolder.getContent()).andReturn(StringUtils.format("\"%s\"", status.toString())).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); - KafkaIndexTask.Status results = client.getStatus(TEST_ID); + Status results = client.getStatus(TEST_ID); verifyAll(); Request request = captured.getValue(); @@ -430,7 +456,11 @@ public void testPause() throws Exception Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).times(2); expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -462,13 +492,25 @@ public void testPauseWithSubsequentGetOffsets() throws Exception .andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn("\"PAUSED\"").times(2) .andReturn("{\"0\":1, \"1\":10}").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); - expect(httpClient.go(capture(captured2), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured2), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); - expect(httpClient.go(capture(captured3), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured3), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); @@ -509,7 +551,11 @@ public void testResume() throws Exception { Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -533,7 +579,11 @@ public void testSetEndOffsets() throws Exception Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -558,7 +608,11 @@ public void testSetEndOffsetsAndResume() throws Exception Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -581,7 +635,11 @@ public void testStop() throws Exception { Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -603,7 +661,11 @@ public void testStopAndPublish() throws Exception { Capture captured = Capture.newInstance(); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ); replayAll(); @@ -626,7 +688,11 @@ public void testStopAsync() throws Exception final int numRequests = TEST_IDS.size(); Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -658,7 +724,11 @@ public void testResumeAsync() throws Exception final int numRequests = TEST_IDS.size(); Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -691,7 +761,11 @@ public void testPauseAsync() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -724,19 +798,23 @@ public void testGetStatusAsync() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn("\"READING\"").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); List expectedUrls = new ArrayList<>(); - List> futures = new ArrayList<>(); + List> futures = new ArrayList<>(); for (String testId : TEST_IDS) { expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "status"))); futures.add(client.getStatusAsync(testId)); } - List responses = Futures.allAsList(futures).get(); + List responses = Futures.allAsList(futures).get(); verifyAll(); List requests = captured.getValues(); @@ -746,7 +824,7 @@ public void testGetStatusAsync() throws Exception for (int i = 0; i < numRequests; i++) { Assert.assertEquals(HttpMethod.GET, requests.get(i).getMethod()); Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); - Assert.assertEquals(KafkaIndexTask.Status.READING, responses.get(i)); + Assert.assertEquals(Status.READING, responses.get(i)); } } @@ -758,7 +836,11 @@ public void testGetStartTimeAsync() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn(String.valueOf(now.getMillis())).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -791,7 +873,11 @@ public void testGetCurrentOffsetsAsync() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -824,7 +910,11 @@ public void testGetEndOffsetsAsync() throws Exception Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -857,7 +947,11 @@ public void testSetEndOffsetsAsync() throws Exception final int numRequests = TEST_IDS.size(); Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); @@ -896,7 +990,11 @@ public void testSetEndOffsetsAsyncWithResume() throws Exception final int numRequests = TEST_IDS.size(); Capture captured = Capture.newInstance(CaptureType.ALL); expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); - expect(httpClient.go(capture(captured), anyObject(FullResponseHandler.class), eq(TEST_HTTP_TIMEOUT))).andReturn( + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( Futures.immediateFuture(responseHolder) ).times(numRequests); replayAll(); diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java index 45c8a93226db..72041eb947d5 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java @@ -69,7 +69,6 @@ import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; import org.apache.druid.indexing.common.task.IndexTaskTest; import org.apache.druid.indexing.common.task.Task; -import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisor; import org.apache.druid.indexing.kafka.test.TestBroker; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; @@ -77,6 +76,9 @@ import org.apache.druid.indexing.overlord.TaskLockbox; import org.apache.druid.indexing.overlord.TaskStorage; import org.apache.druid.indexing.overlord.supervisor.SupervisorManager; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; import org.apache.druid.indexing.test.TestDataSegmentAnnouncer; import org.apache.druid.indexing.test.TestDataSegmentKiller; import org.apache.druid.java.util.common.CompressionUtils; @@ -101,6 +103,7 @@ import org.apache.druid.query.Druids; import org.apache.druid.query.IntervalChunkingQueryRunnerDecorator; import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; import org.apache.druid.query.QueryRunner; import org.apache.druid.query.QueryRunnerFactoryConglomerate; import org.apache.druid.query.QueryToolChest; @@ -174,8 +177,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import static org.apache.druid.query.QueryPlus.wrap; - @RunWith(Parameterized.class) public class KafkaIndexTaskTest { @@ -379,11 +380,11 @@ public void testRunAfterDataInserted() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -407,7 +408,7 @@ public void testRunAfterDataInserted() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -421,11 +422,11 @@ public void testRunBeforeDataInserted() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -437,7 +438,7 @@ public void testRunBeforeDataInserted() throws Exception final ListenableFuture future = runTask(task); // Wait for the task to start reading - while (task.getRunner().getStatus() != KafkaIndexTask.Status.READING) { + while (task.getRunner().getStatus() != Status.READING) { Thread.sleep(10); } @@ -461,7 +462,7 @@ public void testRunBeforeDataInserted() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -489,15 +490,47 @@ public void testIncrementalHandOff() throws Exception Map consumerProps = kafkaServer.consumerProperties(); consumerProps.put("max.poll.records", "1"); - final KafkaPartitions startPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 0L, 1, 0L)); + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 0L, + 1, + 0L + ) + ); // Checkpointing will happen at either checkpoint1 or checkpoint2 depending on ordering // of events fetched across two partitions from Kafka - final KafkaPartitions checkpoint1 = new KafkaPartitions(topic, ImmutableMap.of(0, 5L, 1, 0L)); - final KafkaPartitions checkpoint2 = new KafkaPartitions(topic, ImmutableMap.of(0, 4L, 1, 2L)); - final KafkaPartitions endPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 2L)); + final SeekableStreamPartitions checkpoint1 = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 5L, + 1, + 0L + ) + ); + final SeekableStreamPartitions checkpoint2 = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 4L, + 1, + 2L + ) + ); + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 10L, + 1, + 2L + ) + ); final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, baseSequenceName, startPartitions, @@ -510,12 +543,13 @@ public void testIncrementalHandOff() throws Exception ) ); final ListenableFuture future = runTask(task); - while (task.getRunner().getStatus() != KafkaIndexTask.Status.PAUSED) { + while (task.getRunner().getStatus() != Status.PAUSED) { Thread.sleep(10); } final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); - Assert.assertTrue(checkpoint1.getPartitionOffsetMap().equals(currentOffsets) || checkpoint2.getPartitionOffsetMap() - .equals(currentOffsets)); + Assert.assertTrue(checkpoint1.getPartitionSequenceNumberMap().equals(currentOffsets) + || checkpoint2.getPartitionSequenceNumberMap() + .equals(currentOffsets)); task.getRunner().setEndOffsets(currentOffsets, false); Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); @@ -526,7 +560,7 @@ public void testIncrementalHandOff() throws Exception DATA_SCHEMA.getDataSource(), 0, new KafkaDataSourceMetadata(startPartitions), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, currentOffsets)) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, currentOffsets)) ) ) ); @@ -546,7 +580,7 @@ public void testIncrementalHandOff() throws Exception SegmentDescriptor desc7 = SD(task, "2013/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4, desc5, desc6, desc7), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 2L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 2L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -584,14 +618,46 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception Map consumerProps = kafkaServer.consumerProperties(); consumerProps.put("max.poll.records", "1"); - final KafkaPartitions startPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 0L, 1, 0L)); - final KafkaPartitions checkpoint1 = new KafkaPartitions(topic, ImmutableMap.of(0, 3L, 1, 0L)); - final KafkaPartitions checkpoint2 = new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 0L)); + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 0L, + 1, + 0L + ) + ); + final SeekableStreamPartitions checkpoint1 = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 3L, + 1, + 0L + ) + ); + final SeekableStreamPartitions checkpoint2 = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 10L, + 1, + 0L + ) + ); - final KafkaPartitions endPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 2L)); + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 10L, + 1, + 2L + ) + ); final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, baseSequenceName, startPartitions, @@ -604,15 +670,15 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception ) ); final ListenableFuture future = runTask(task); - while (task.getRunner().getStatus() != KafkaIndexTask.Status.PAUSED) { + while (task.getRunner().getStatus() != Status.PAUSED) { Thread.sleep(10); } final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); - Assert.assertTrue(checkpoint1.getPartitionOffsetMap().equals(currentOffsets)); + Assert.assertTrue(checkpoint1.getPartitionSequenceNumberMap().equals(currentOffsets)); task.getRunner().setEndOffsets(currentOffsets, false); - while (task.getRunner().getStatus() != KafkaIndexTask.Status.PAUSED) { + while (task.getRunner().getStatus() != Status.PAUSED) { Thread.sleep(10); } @@ -622,7 +688,7 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception } final Map nextOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); - Assert.assertTrue(checkpoint2.getPartitionOffsetMap().equals(nextOffsets)); + Assert.assertTrue(checkpoint2.getPartitionSequenceNumberMap().equals(nextOffsets)); task.getRunner().setEndOffsets(nextOffsets, false); Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); @@ -634,7 +700,7 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception DATA_SCHEMA.getDataSource(), 0, new KafkaDataSourceMetadata(startPartitions), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, currentOffsets)) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, currentOffsets)) ) ) ); @@ -643,8 +709,8 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception Objects.hash( DATA_SCHEMA.getDataSource(), 0, - new KafkaDataSourceMetadata(new KafkaPartitions(topic, currentOffsets)), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, nextOffsets)) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, currentOffsets)), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, nextOffsets)) ) ) ); @@ -664,7 +730,7 @@ public void testIncrementalHandOffMaxTotalRows() throws Exception SegmentDescriptor desc7 = SD(task, "2013/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4, desc5, desc6, desc7), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 2L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 2L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -701,13 +767,37 @@ public void testTimeBasedIncrementalHandOff() throws Exception Map consumerProps = kafkaServer.consumerProperties(); consumerProps.put("max.poll.records", "1"); - final KafkaPartitions startPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 0L, 1, 0L)); + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 0L, + 1, + 0L + ) + ); // Checkpointing will happen at checkpoint - final KafkaPartitions checkpoint = new KafkaPartitions(topic, ImmutableMap.of(0, 1L, 1, 0L)); - final KafkaPartitions endPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 2L, 1, 0L)); + final SeekableStreamPartitions checkpoint = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 1L, + 1, + 0L + ) + ); + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of( + 0, + 2L, + 1, + 0L + ) + ); final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, baseSequenceName, startPartitions, @@ -722,11 +812,11 @@ public void testTimeBasedIncrementalHandOff() throws Exception final ListenableFuture future = runTask(task); // task will pause for checkpointing - while (task.getRunner().getStatus() != KafkaIndexTask.Status.PAUSED) { + while (task.getRunner().getStatus() != Status.PAUSED) { Thread.sleep(10); } final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); - Assert.assertTrue(checkpoint.getPartitionOffsetMap().equals(currentOffsets)); + Assert.assertTrue(checkpoint.getPartitionSequenceNumberMap().equals(currentOffsets)); task.getRunner().setEndOffsets(currentOffsets, false); Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); @@ -737,7 +827,10 @@ public void testTimeBasedIncrementalHandOff() throws Exception DATA_SCHEMA.getDataSource(), 0, new KafkaDataSourceMetadata(startPartitions), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, checkpoint.getPartitionOffsetMap())) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( + topic, + checkpoint.getPartitionSequenceNumberMap() + )) ) ) ); @@ -752,7 +845,7 @@ public void testTimeBasedIncrementalHandOff() throws Exception SegmentDescriptor desc2 = SD(task, "2009/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 2L, 1, 0L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L, 1, 0L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -791,13 +884,22 @@ public void testIncrementalHandOffReadsThroughEndOffsets() throws Exception Map consumerProps = kafkaServer.consumerProperties(); consumerProps.put("max.poll.records", "1"); - final KafkaPartitions startPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 0L)); - final KafkaPartitions checkpoint1 = new KafkaPartitions(topic, ImmutableMap.of(0, 5L)); - final KafkaPartitions endPartitions = new KafkaPartitions(topic, ImmutableMap.of(0, 7L)); + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of(0, 0L) + ); + final SeekableStreamPartitions checkpoint1 = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of(0, 5L) + ); + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + topic, + ImmutableMap.of(0, 7L) + ); final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, baseSequenceName, startPartitions, @@ -810,11 +912,11 @@ public void testIncrementalHandOffReadsThroughEndOffsets() throws Exception ) ); final ListenableFuture future = runTask(task); - while (task.getRunner().getStatus() != KafkaIndexTask.Status.PAUSED) { + while (task.getRunner().getStatus() != Status.PAUSED) { Thread.sleep(10); } final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); - Assert.assertTrue(checkpoint1.getPartitionOffsetMap().equals(currentOffsets)); + Assert.assertTrue(checkpoint1.getPartitionSequenceNumberMap().equals(currentOffsets)); // actual checkpoint offset is 5, but simulating behavior of publishing set end offset call, to ensure this task // will continue reading through the end offset of the checkpointed sequence @@ -833,11 +935,11 @@ public void testRunWithMinimumMessageTime() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, DateTimes.of("2010"), @@ -849,7 +951,7 @@ public void testRunWithMinimumMessageTime() throws Exception final ListenableFuture future = runTask(task); // Wait for the task to start reading - while (task.getRunner().getStatus() != KafkaIndexTask.Status.READING) { + while (task.getRunner().getStatus() != Status.READING) { Thread.sleep(10); } @@ -873,7 +975,7 @@ public void testRunWithMinimumMessageTime() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -887,11 +989,11 @@ public void testRunWithMaximumMessageTime() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -903,7 +1005,7 @@ public void testRunWithMaximumMessageTime() throws Exception final ListenableFuture future = runTask(task); // Wait for the task to start reading - while (task.getRunner().getStatus() != KafkaIndexTask.Status.READING) { + while (task.getRunner().getStatus() != Status.READING) { Thread.sleep(10); } @@ -928,7 +1030,7 @@ public void testRunWithMaximumMessageTime() throws Exception SegmentDescriptor desc3 = SD(task, "2010/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -951,11 +1053,11 @@ public void testRunWithTransformSpec() throws Exception ) ) ), - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -967,7 +1069,7 @@ public void testRunWithTransformSpec() throws Exception final ListenableFuture future = runTask(task); // Wait for the task to start reading - while (task.getRunner().getStatus() != KafkaIndexTask.Status.READING) { + while (task.getRunner().getStatus() != Status.READING) { Thread.sleep(10); } @@ -990,7 +1092,7 @@ public void testRunWithTransformSpec() throws Exception SegmentDescriptor desc1 = SD(task, "2009/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1011,11 +1113,11 @@ public void testRunOnNothing() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), kafkaServer.consumerProperties(), true, null, @@ -1052,11 +1154,11 @@ public void testHandoffConditionTimeoutWhenHandoffOccurs() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1080,7 +1182,7 @@ public void testHandoffConditionTimeoutWhenHandoffOccurs() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1104,11 +1206,11 @@ public void testHandoffConditionTimeoutWhenHandoffDoesNotOccur() throws Exceptio final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1132,7 +1234,7 @@ public void testHandoffConditionTimeoutWhenHandoffDoesNotOccur() throws Exceptio SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1159,11 +1261,11 @@ public void testReportParseExceptions() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 7L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 7L)), kafkaServer.consumerProperties(), true, null, @@ -1203,11 +1305,11 @@ public void testMultipleParseExceptionsSuccess() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 13L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 13L)), kafkaServer.consumerProperties(), true, null, @@ -1237,7 +1339,7 @@ public void testMultipleParseExceptionsSuccess() throws Exception SegmentDescriptor desc4 = SD(task, "2049/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 13L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 13L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1285,11 +1387,11 @@ public void testMultipleParseExceptionsFailure() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 10L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L)), kafkaServer.consumerProperties(), true, null, @@ -1345,11 +1447,11 @@ public void testRunReplicas() throws Exception { final KafkaIndexTask task1 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1359,11 +1461,11 @@ public void testRunReplicas() throws Exception ); final KafkaIndexTask task2 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1399,7 +1501,7 @@ public void testRunReplicas() throws Exception SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1413,11 +1515,11 @@ public void testRunConflicting() throws Exception { final KafkaIndexTask task1 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1427,11 +1529,11 @@ public void testRunConflicting() throws Exception ); final KafkaIndexTask task2 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 1, "sequence1", - new KafkaPartitions(topic, ImmutableMap.of(0, 3L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 10L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 3L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L)), kafkaServer.consumerProperties(), true, null, @@ -1468,7 +1570,7 @@ public void testRunConflicting() throws Exception SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1482,11 +1584,11 @@ public void testRunConflictingWithoutTransactions() throws Exception { final KafkaIndexTask task1 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), false, null, @@ -1496,11 +1598,11 @@ public void testRunConflictingWithoutTransactions() throws Exception ); final KafkaIndexTask task2 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 1, "sequence1", - new KafkaPartitions(topic, ImmutableMap.of(0, 3L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 10L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 3L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L)), kafkaServer.consumerProperties(), false, null, @@ -1556,11 +1658,11 @@ public void testRunOneTaskTwoPartitions() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L, 1, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L, 1, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L, 1, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L, 1, 2L)), kafkaServer.consumerProperties(), true, null, @@ -1597,7 +1699,7 @@ public void testRunOneTaskTwoPartitions() throws Exception ? ImmutableSet.of(desc1, desc2, desc4) : ImmutableSet.of(desc1, desc2, desc3, desc4), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L, 1, 2L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L, 1, 2L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1621,11 +1723,11 @@ public void testRunTwoTasksTwoPartitions() throws Exception { final KafkaIndexTask task1 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1635,11 +1737,11 @@ public void testRunTwoTasksTwoPartitions() throws Exception ); final KafkaIndexTask task2 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 1, "sequence1", - new KafkaPartitions(topic, ImmutableMap.of(1, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(1, 1L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(1, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(1, 1L)), kafkaServer.consumerProperties(), true, null, @@ -1676,7 +1778,7 @@ public void testRunTwoTasksTwoPartitions() throws Exception SegmentDescriptor desc3 = SD(task2, "2012/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L, 1, 1L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L, 1, 1L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1691,11 +1793,11 @@ public void testRestore() throws Exception { final KafkaIndexTask task1 = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1728,11 +1830,11 @@ public void testRestore() throws Exception // Start a new task final KafkaIndexTask task2 = createTask( task1.getId(), - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1766,7 +1868,7 @@ public void testRestore() throws Exception SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1780,11 +1882,11 @@ public void testRunWithPauseAndResume() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1808,7 +1910,7 @@ public void testRunWithPauseAndResume() throws Exception } Assert.assertEquals(2, countEvents(task)); - Assert.assertEquals(KafkaIndexTask.Status.READING, task.getRunner().getStatus()); + Assert.assertEquals(Status.READING, task.getRunner().getStatus()); Map currentOffsets = objectMapper.readValue( task.getRunner().pause().getEntity().toString(), @@ -1816,7 +1918,7 @@ public void testRunWithPauseAndResume() throws Exception { } ); - Assert.assertEquals(KafkaIndexTask.Status.PAUSED, task.getRunner().getStatus()); + Assert.assertEquals(Status.PAUSED, task.getRunner().getStatus()); // Insert remaining data try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { @@ -1850,7 +1952,7 @@ public void testRunWithPauseAndResume() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -1864,11 +1966,11 @@ public void testRunWithOffsetOutOfRangeExceptionAndPause() throws Exception { final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 2L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 2L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1879,13 +1981,13 @@ public void testRunWithOffsetOutOfRangeExceptionAndPause() throws Exception runTask(task); - while (!task.getRunner().getStatus().equals(KafkaIndexTask.Status.READING)) { + while (!task.getRunner().getStatus().equals(Status.READING)) { Thread.sleep(2000); } task.getRunner().pause(); - while (!task.getRunner().getStatus().equals(KafkaIndexTask.Status.PAUSED)) { + while (!task.getRunner().getStatus().equals(Status.PAUSED)) { Thread.sleep(25); } } @@ -1903,11 +2005,11 @@ public void testRunWithOffsetOutOfRangeExceptionAndNextOffsetGreaterThanLeastAva final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", - new KafkaPartitions(topic, ImmutableMap.of(0, 200L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 500L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 200L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 500L)), kafkaServer.consumerProperties(), true, null, @@ -1918,12 +2020,12 @@ public void testRunWithOffsetOutOfRangeExceptionAndNextOffsetGreaterThanLeastAva runTask(task); - while (!task.getRunner().getStatus().equals(KafkaIndexTask.Status.READING)) { + while (!task.getRunner().getStatus().equals(Status.READING)) { Thread.sleep(20); } for (int i = 0; i < 5; i++) { - Assert.assertEquals(task.getRunner().getStatus(), KafkaIndexTask.Status.READING); + Assert.assertEquals(task.getRunner().getStatus(), Status.READING); // Offset should not be reset Assert.assertTrue(task.getRunner().getCurrentOffsets().get(0) == 200L); } @@ -1956,12 +2058,12 @@ public void testRunContextSequenceAheadOfStartingOffsets() throws Exception final KafkaIndexTask task = createTask( null, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( 0, "sequence0", // task should ignore these and use sequence info sent in the context - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, 5L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L)), kafkaServer.consumerProperties(), true, null, @@ -1986,7 +2088,7 @@ public void testRunContextSequenceAheadOfStartingOffsets() throws Exception SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); Assert.assertEquals( - new KafkaDataSourceMetadata(new KafkaPartitions(topic, ImmutableMap.of(0, 5L))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 5L))), metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) ); @@ -2042,7 +2144,7 @@ public boolean apply(TaskLock lock) private KafkaIndexTask createTask( final String taskId, - final KafkaIOConfig ioConfig + final KafkaIndexTaskIOConfig ioConfig ) { return createTask(taskId, DATA_SCHEMA, ioConfig); @@ -2050,7 +2152,7 @@ private KafkaIndexTask createTask( private KafkaIndexTask createTask( final String taskId, - final KafkaIOConfig ioConfig, + final KafkaIndexTaskIOConfig ioConfig, final Map context ) { @@ -2060,10 +2162,10 @@ private KafkaIndexTask createTask( private KafkaIndexTask createTask( final String taskId, final DataSchema dataSchema, - final KafkaIOConfig ioConfig + final KafkaIndexTaskIOConfig ioConfig ) { - final KafkaTuningConfig tuningConfig = new KafkaTuningConfig( + final KafkaIndexTaskTuningConfig tuningConfig = new KafkaIndexTaskTuningConfig( 1000, null, maxRowsPerSegment, @@ -2083,7 +2185,10 @@ private KafkaIndexTask createTask( maxSavedParseExceptions ); final Map context = isIncrementalHandoffSupported - ? ImmutableMap.of(KafkaSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED, true) + ? ImmutableMap.of( + SeekableStreamSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED, + true + ) : null; final KafkaIndexTask task = new KafkaIndexTask( taskId, @@ -2105,11 +2210,11 @@ private KafkaIndexTask createTask( private KafkaIndexTask createTask( final String taskId, final DataSchema dataSchema, - final KafkaIOConfig ioConfig, + final KafkaIndexTaskIOConfig ioConfig, final Map context ) { - final KafkaTuningConfig tuningConfig = new KafkaTuningConfig( + final KafkaIndexTaskTuningConfig tuningConfig = new KafkaIndexTaskTuningConfig( 1000, null, maxRowsPerSegment, @@ -2129,7 +2234,7 @@ private KafkaIndexTask createTask( maxSavedParseExceptions ); if (isIncrementalHandoffSupported) { - context.put(KafkaSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED, true); + context.put(SeekableStreamSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED, true); } final KafkaIndexTask task = new KafkaIndexTask( @@ -2401,7 +2506,8 @@ private List readSegmentColumn(final String column, final SegmentDescrip ); IndexIO indexIO = new TestUtils().getTestIndexIO(); QueryableIndex index = indexIO.loadIndex(outputLocation); - DictionaryEncodedColumn theColumn = (DictionaryEncodedColumn) index.getColumnHolder(column).getColumn(); + DictionaryEncodedColumn theColumn = (DictionaryEncodedColumn) index.getColumnHolder(column) + .getColumn(); List values = new ArrayList<>(); for (int i = 0; i < theColumn.length(); i++) { int id = theColumn.getSingleValueRow(i); @@ -2425,7 +2531,7 @@ public long countEvents(final Task task) .build(); List> results = - task.getQueryRunner(query).run(wrap(query), ImmutableMap.of()).toList(); + task.getQueryRunner(query).run(QueryPlus.wrap(query), ImmutableMap.of()).toList(); return results.isEmpty() ? 0L : DimensionHandlerUtils.nullToZero(results.get(0).getValue().getLongMetric("rows")); } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaTuningConfigTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java similarity index 74% rename from extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaTuningConfigTest.java rename to extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java index 69cf186a3c3b..abcf07f6c10b 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaTuningConfigTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.Module; import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorTuningConfig; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.segment.IndexSpec; import org.apache.druid.segment.indexing.TuningConfig; @@ -30,11 +31,11 @@ import java.io.File; -public class KafkaTuningConfigTest +public class KafkaIndexTaskTuningConfigTest { private final ObjectMapper mapper; - public KafkaTuningConfigTest() + public KafkaIndexTaskTuningConfigTest() { mapper = new DefaultObjectMapper(); mapper.registerModules((Iterable) new KafkaIndexTaskModule().getJacksonModules()); @@ -45,7 +46,7 @@ public void testSerdeWithDefaults() throws Exception { String jsonStr = "{\"type\": \"kafka\"}"; - KafkaTuningConfig config = (KafkaTuningConfig) mapper.readValue( + KafkaIndexTaskTuningConfig config = (KafkaIndexTaskTuningConfig) mapper.readValue( mapper.writeValueAsString( mapper.readValue( jsonStr, @@ -81,7 +82,7 @@ public void testSerdeWithNonDefaults() throws Exception + " \"handoffConditionTimeout\": 100\n" + "}"; - KafkaTuningConfig config = (KafkaTuningConfig) mapper.readValue( + KafkaIndexTaskTuningConfig config = (KafkaIndexTaskTuningConfig) mapper.readValue( mapper.writeValueAsString( mapper.readValue( jsonStr, @@ -97,15 +98,15 @@ public void testSerdeWithNonDefaults() throws Exception Assert.assertNotEquals(null, config.getMaxTotalRows()); Assert.assertEquals(1000, config.getMaxTotalRows().longValue()); Assert.assertEquals(new Period("PT1H"), config.getIntermediatePersistPeriod()); - Assert.assertEquals(0, config.getMaxPendingPersists()); + Assert.assertEquals(100, config.getMaxPendingPersists()); Assert.assertEquals(true, config.isReportParseExceptions()); Assert.assertEquals(100, config.getHandoffConditionTimeout()); } @Test - public void testCopyOf() + public void testConvert() { - KafkaTuningConfig original = new KafkaTuningConfig( + KafkaSupervisorTuningConfig original = new KafkaSupervisorTuningConfig( 1, null, 2, @@ -122,9 +123,15 @@ public void testCopyOf() null, null, null, + null, + null, + null, + null, + null, + null, null ); - KafkaTuningConfig copy = KafkaTuningConfig.copyOf(original); + KafkaIndexTaskTuningConfig copy = (KafkaIndexTaskTuningConfig) original.convertToTaskTuningConfig(); Assert.assertEquals(1, copy.getMaxRowsInMemory()); Assert.assertEquals(2, copy.getMaxRowsPerSegment()); @@ -132,9 +139,32 @@ public void testCopyOf() Assert.assertEquals(10L, copy.getMaxTotalRows().longValue()); Assert.assertEquals(new Period("PT3S"), copy.getIntermediatePersistPeriod()); Assert.assertEquals(new File("/tmp/xxx"), copy.getBasePersistDirectory()); - Assert.assertEquals(0, copy.getMaxPendingPersists()); + Assert.assertEquals(4, copy.getMaxPendingPersists()); Assert.assertEquals(new IndexSpec(), copy.getIndexSpec()); Assert.assertEquals(true, copy.isReportParseExceptions()); Assert.assertEquals(5L, copy.getHandoffConditionTimeout()); } + + private static KafkaIndexTaskTuningConfig copy(KafkaIndexTaskTuningConfig config) + { + return new KafkaIndexTaskTuningConfig( + config.getMaxRowsInMemory(), + config.getMaxBytesInMemory(), + config.getMaxRowsPerSegment(), + config.getMaxTotalRows(), + config.getIntermediatePersistPeriod(), + config.getBasePersistDirectory(), + 0, + config.getIndexSpec(), + true, + config.isReportParseExceptions(), + config.getHandoffConditionTimeout(), + config.isResetOffsetAutomatically(), + config.getSegmentWriteOutMediumFactory(), + config.getIntermediateHandoffPeriod(), + config.isLogParseExceptions(), + config.getMaxParseExceptions(), + config.getMaxSavedParseExceptions() + ); + } } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java new file mode 100644 index 000000000000..2f445aa99431 --- /dev/null +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java @@ -0,0 +1,437 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kafka; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.curator.test.TestingCluster; +import org.apache.druid.indexing.kafka.test.TestBroker; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.TestHelper; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; + +public class KafkaRecordSupplierTest +{ + private static final Logger log = new Logger(KafkaRecordSupplierTest.class); + private static String topic = "topic"; + private static long poll_timeout_millis = 1000; + private static int pollRetry = 5; + private static int topicPosFix = 0; + private static final ObjectMapper objectMapper = TestHelper.makeJsonMapper(); + + private static TestingCluster zkServer; + private static TestBroker kafkaServer; + + private List> records; + + + private static List> generateRecords(String topic) + { + return ImmutableList.of( + new ProducerRecord<>(topic, 0, null, JB("2008", "a", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("2009", "b", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("2010", "c", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("2011", "d", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("2011", "e", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("246140482-04-24T15:36:27.903Z", "x", "z", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, StringUtils.toUtf8("unparseable")), + new ProducerRecord<>(topic, 0, null, StringUtils.toUtf8("unparseable2")), + new ProducerRecord<>(topic, 0, null, null), + new ProducerRecord<>(topic, 0, null, JB("2013", "f", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 0, null, JB("2049", "f", "y", "notanumber", "20.0", "1.0")), + new ProducerRecord<>(topic, 1, null, JB("2049", "f", "y", "10", "notanumber", "1.0")), + new ProducerRecord<>(topic, 1, null, JB("2049", "f", "y", "10", "20.0", "notanumber")), + new ProducerRecord<>(topic, 1, null, JB("2012", "g", "y", "10", "20.0", "1.0")), + new ProducerRecord<>(topic, 1, null, JB("2011", "h", "y", "10", "20.0", "1.0")) + ); + } + + private static byte[] JB(String timestamp, String dim1, String dim2, String dimLong, String dimFloat, String met1) + { + try { + return new ObjectMapper().writeValueAsBytes( + ImmutableMap.builder() + .put("timestamp", timestamp) + .put("dim1", dim1) + .put("dim2", dim2) + .put("dimLong", dimLong) + .put("dimFloat", dimFloat) + .put("met1", met1) + .build() + ); + } + catch (Exception e) { + throw Throwables.propagate(e); + } + } + + private static String getTopicName() + { + return "topic-" + topicPosFix++; + } + + private List> createOrderedPartitionableRecords() + { + Map partitionToOffset = new HashMap<>(); + return records.stream().map(r -> { + long offset = 0; + if (partitionToOffset.containsKey(r.partition())) { + offset = partitionToOffset.get(r.partition()); + partitionToOffset.put(r.partition(), offset + 1); + } else { + partitionToOffset.put(r.partition(), 1L); + } + return new OrderedPartitionableRecord<>( + topic, + r.partition(), + offset, + r.value() == null ? null : Collections.singletonList(r.value()) + ); + }).collect(Collectors.toList()); + } + + @BeforeClass + public static void setupClass() throws Exception + { + zkServer = new TestingCluster(1); + zkServer.start(); + + kafkaServer = new TestBroker( + zkServer.getConnectString(), + null, + 1, + ImmutableMap.of("num.partitions", "2") + ); + kafkaServer.start(); + + } + + @Before + public void setupTest() + { + topic = getTopicName(); + records = generateRecords(topic); + } + + @AfterClass + public static void tearDownClass() throws Exception + { + kafkaServer.close(); + kafkaServer = null; + + zkServer.stop(); + zkServer = null; + } + + @Test + public void testSupplierSetup() throws ExecutionException, InterruptedException + { + + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + Assert.assertTrue(recordSupplier.getAssignment().isEmpty()); + + recordSupplier.assign(partitions); + + Assert.assertEquals(partitions, recordSupplier.getAssignment()); + Assert.assertEquals(ImmutableSet.of(0, 1), recordSupplier.getPartitionIds(topic)); + + recordSupplier.close(); + } + + @Test + public void testPoll() throws InterruptedException, ExecutionException + { + + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + recordSupplier.seekToEarliest(partitions); + + List> initialRecords = new ArrayList<>(createOrderedPartitionableRecords()); + + List> polledRecords = recordSupplier.poll(poll_timeout_millis); + for (int i = 0; polledRecords.size() != initialRecords.size() && i < pollRetry; i++) { + polledRecords.addAll(recordSupplier.poll(poll_timeout_millis)); + Thread.sleep(200); + } + + Assert.assertEquals(partitions, recordSupplier.getAssignment()); + Assert.assertEquals(initialRecords.size(), polledRecords.size()); + Assert.assertTrue(initialRecords.containsAll(polledRecords)); + + recordSupplier.close(); + } + + + @Test + public void testPollAfterMoreDataAdded() throws InterruptedException, ExecutionException + { + // Insert data + + KafkaProducer producer = kafkaServer.newProducer(); + for (ProducerRecord record : records.subList(0, 13)) { + producer.send(record).get(); + } + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + recordSupplier.seekToEarliest(partitions); + + List> polledRecords = recordSupplier.poll(poll_timeout_millis); + for (int i = 0; polledRecords.size() != 13 && i < pollRetry; i++) { + polledRecords.addAll(recordSupplier.poll(poll_timeout_millis)); + Thread.sleep(200); + } + + // Insert data + for (ProducerRecord rec : records.subList(13, 15)) { + producer.send(rec).get(); + } + + + for (int i = 0; polledRecords.size() != records.size() && i < pollRetry; i++) { + polledRecords.addAll(recordSupplier.poll(poll_timeout_millis)); + Thread.sleep(200); + } + + List> initialRecords = createOrderedPartitionableRecords(); + + Assert.assertEquals(records.size(), polledRecords.size()); + Assert.assertTrue(initialRecords.containsAll(polledRecords)); + + + recordSupplier.close(); + } + + @Test + public void testSeek() throws InterruptedException, ExecutionException + { + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + StreamPartition partition0 = StreamPartition.of(topic, 0); + StreamPartition partition1 = StreamPartition.of(topic, 1); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + + Assert.assertEquals(0L, (long) recordSupplier.getEarliestSequenceNumber(partition0)); + Assert.assertEquals(0L, (long) recordSupplier.getEarliestSequenceNumber(partition1)); + + recordSupplier.seek(partition0, 2L); + recordSupplier.seek(partition1, 2L); + + List> initialRecords = createOrderedPartitionableRecords(); + + List> polledRecords = recordSupplier.poll(poll_timeout_millis); + for (int i = 0; polledRecords.size() != 11 && i < pollRetry; i++) { + polledRecords.addAll(recordSupplier.poll(poll_timeout_millis)); + Thread.sleep(200); + } + + + Assert.assertEquals(11, polledRecords.size()); + Assert.assertTrue(initialRecords.containsAll(polledRecords)); + + + recordSupplier.close(); + + } + + @Test + public void testSeekToLatest() throws InterruptedException, ExecutionException + { + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + StreamPartition partition0 = StreamPartition.of(topic, 0); + StreamPartition partition1 = StreamPartition.of(topic, 1); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + + Assert.assertEquals(0L, (long) recordSupplier.getEarliestSequenceNumber(partition0)); + Assert.assertEquals(0L, (long) recordSupplier.getEarliestSequenceNumber(partition1)); + + recordSupplier.seekToLatest(partitions); + List> polledRecords = recordSupplier.poll(poll_timeout_millis); + + Assert.assertEquals(Collections.emptyList(), polledRecords); + recordSupplier.close(); + } + + @Test(expected = IllegalStateException.class) + public void testSeekUnassigned() throws InterruptedException, ExecutionException + { + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + StreamPartition partition0 = StreamPartition.of(topic, 0); + StreamPartition partition1 = StreamPartition.of(topic, 1); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + + Assert.assertEquals(0, (long) recordSupplier.getEarliestSequenceNumber(partition0)); + + recordSupplier.seekToEarliest(Collections.singleton(partition1)); + + recordSupplier.close(); + } + + @Test + public void testPosition() throws ExecutionException, InterruptedException + { + // Insert data + try (final KafkaProducer kafkaProducer = kafkaServer.newProducer()) { + for (ProducerRecord record : records) { + kafkaProducer.send(record).get(); + } + } + + StreamPartition partition0 = StreamPartition.of(topic, 0); + StreamPartition partition1 = StreamPartition.of(topic, 1); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(topic, 0), + StreamPartition.of(topic, 1) + ); + + KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier( + kafkaServer.consumerProperties(), objectMapper); + + recordSupplier.assign(partitions); + + Assert.assertEquals(0L, (long) recordSupplier.getPosition(partition0)); + Assert.assertEquals(0L, (long) recordSupplier.getPosition(partition1)); + + recordSupplier.seek(partition0, 4L); + recordSupplier.seek(partition1, 5L); + + Assert.assertEquals(4L, (long) recordSupplier.getPosition(partition0)); + Assert.assertEquals(5L, (long) recordSupplier.getPosition(partition1)); + + recordSupplier.seekToEarliest(Collections.singleton(partition0)); + Assert.assertEquals(0L, (long) recordSupplier.getPosition(partition0)); + + recordSupplier.seekToLatest(Collections.singleton(partition0)); + Assert.assertEquals(11L, (long) recordSupplier.getPosition(partition0)); + + long prevPos = recordSupplier.getPosition(partition0); + recordSupplier.getEarliestSequenceNumber(partition0); + Assert.assertEquals(prevPos, (long) recordSupplier.getPosition(partition0)); + + recordSupplier.getLatestSequenceNumber(partition0); + Assert.assertEquals(prevPos, (long) recordSupplier.getPosition(partition0)); + + + recordSupplier.close(); + } + + +} diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java index e64015adba58..cbf25def6b4f 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java @@ -23,8 +23,8 @@ import com.fasterxml.jackson.databind.Module; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableMap; -import org.apache.druid.indexing.kafka.KafkaIndexTask; import org.apache.druid.indexing.kafka.KafkaIndexTaskModule; +import org.apache.druid.indexing.kafka.KafkaRecordSupplier; import org.apache.druid.jackson.DefaultObjectMapper; import org.hamcrest.CoreMatchers; import org.joda.time.Duration; @@ -136,7 +136,7 @@ public void testSerdeForConsumerPropertiesWithPasswords() throws Exception KafkaSupervisorIOConfig config = mapper.readValue(jsonStr, KafkaSupervisorIOConfig.class); Properties props = new Properties(); - KafkaIndexTask.addConsumerPropertiesFromConfig(props, mapper, config.getConsumerProperties()); + KafkaRecordSupplier.addConsumerPropertiesFromConfig(props, mapper, config.getConsumerProperties()); Assert.assertEquals("my-topic", config.getTopic()); Assert.assertEquals("localhost:9092", props.getProperty("bootstrap.servers")); @@ -187,4 +187,5 @@ public void testBootstrapServersRequired() throws Exception exception.expectMessage(CoreMatchers.containsString("bootstrap.servers")); mapper.readValue(jsonStr, KafkaSupervisorIOConfig.class); } + } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java index f08212aa309f..427b9d1742de 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java @@ -45,12 +45,10 @@ import org.apache.druid.indexing.common.task.RealtimeIndexTask; import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.kafka.KafkaDataSourceMetadata; -import org.apache.druid.indexing.kafka.KafkaIOConfig; import org.apache.druid.indexing.kafka.KafkaIndexTask; import org.apache.druid.indexing.kafka.KafkaIndexTaskClient; import org.apache.druid.indexing.kafka.KafkaIndexTaskClientFactory; -import org.apache.druid.indexing.kafka.KafkaPartitions; -import org.apache.druid.indexing.kafka.KafkaTuningConfig; +import org.apache.druid.indexing.kafka.KafkaIndexTaskIOConfig; import org.apache.druid.indexing.kafka.test.TestBroker; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; @@ -61,6 +59,9 @@ import org.apache.druid.indexing.overlord.TaskRunnerWorkItem; import org.apache.druid.indexing.overlord.TaskStorage; import org.apache.druid.indexing.overlord.supervisor.SupervisorReport; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.supervisor.TaskReportData; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; @@ -112,15 +113,11 @@ import java.util.concurrent.Executor; import java.util.concurrent.TimeoutException; -import static org.easymock.EasyMock.anyBoolean; import static org.easymock.EasyMock.anyObject; -import static org.easymock.EasyMock.anyString; import static org.easymock.EasyMock.capture; -import static org.easymock.EasyMock.eq; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.expectLastCall; import static org.easymock.EasyMock.replay; -import static org.easymock.EasyMock.reset; @RunWith(Parameterized.class) public class KafkaSupervisorTest extends EasyMockSupport @@ -141,6 +138,7 @@ public class KafkaSupervisorTest extends EasyMockSupport private static int topicPostfix; private static ZkUtils zkUtils; + private final int numThreads; private KafkaSupervisor supervisor; @@ -282,9 +280,9 @@ public void testNoInitialState() throws Exception KafkaIndexTask task = captured.getValue(); Assert.assertEquals(dataSchema, task.getDataSchema()); - Assert.assertEquals(KafkaTuningConfig.copyOf(tuningConfig), task.getTuningConfig()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), task.getTuningConfig()); - KafkaIOConfig taskConfig = task.getIOConfig(); + KafkaIndexTaskIOConfig taskConfig = task.getIOConfig(); Assert.assertEquals(kafkaHost, taskConfig.getConsumerProperties().get("bootstrap.servers")); Assert.assertEquals("myCustomValue", taskConfig.getConsumerProperties().get("myCustomKey")); Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); @@ -293,15 +291,15 @@ public void testNoInitialState() throws Exception Assert.assertFalse("maximumMessageTime", taskConfig.getMaximumMessageTime().isPresent()); Assert.assertFalse("skipOffsetGaps", taskConfig.isSkipOffsetGaps()); - Assert.assertEquals(topic, taskConfig.getStartPartitions().getTopic()); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, taskConfig.getStartPartitions().getStream()); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); - Assert.assertEquals(topic, taskConfig.getEndPartitions().getTopic()); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, taskConfig.getEndPartitions().getStream()); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(2)); } @Test @@ -328,7 +326,7 @@ public void testSkipOffsetGaps() throws Exception verifyAll(); KafkaIndexTask task = captured.getValue(); - KafkaIOConfig taskConfig = task.getIOConfig(); + KafkaIndexTaskIOConfig taskConfig = task.getIOConfig(); Assert.assertTrue("skipOffsetGaps", taskConfig.isSkipOffsetGaps()); } @@ -356,18 +354,27 @@ public void testMultiTask() throws Exception verifyAll(); KafkaIndexTask task1 = captured.getValues().get(0); - Assert.assertEquals(2, task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(2, task1.getIOConfig().getEndPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(Long.MAX_VALUE, (long) task1.getIOConfig().getEndPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(2)); - Assert.assertEquals(Long.MAX_VALUE, (long) task1.getIOConfig().getEndPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(2, task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(2, task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals( + Long.MAX_VALUE, + (long) task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(0) + ); + Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(2)); + Assert.assertEquals( + Long.MAX_VALUE, + (long) task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(2) + ); KafkaIndexTask task2 = captured.getValues().get(1); - Assert.assertEquals(1, task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(1, task2.getIOConfig().getEndPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(Long.MAX_VALUE, (long) task2.getIOConfig().getEndPartitions().getPartitionOffsetMap().get(1)); + Assert.assertEquals(1, task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(1, task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals( + Long.MAX_VALUE, + (long) task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(1) + ); } @Test @@ -393,18 +400,18 @@ public void testReplicas() throws Exception verifyAll(); KafkaIndexTask task1 = captured.getValues().get(0); - Assert.assertEquals(3, task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(3, task1.getIOConfig().getEndPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(3, task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(3, task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(0L, (long) task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(2)); KafkaIndexTask task2 = captured.getValues().get(1); - Assert.assertEquals(3, task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(3, task2.getIOConfig().getEndPartitions().getPartitionOffsetMap().size()); - Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(3, task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(3, task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(0L, (long) task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(2)); } @Test @@ -511,9 +518,9 @@ public void testLatestOffset() throws Exception verifyAll(); KafkaIndexTask task = captured.getValue(); - Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(1100L, (long) task.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(2)); } @Test @@ -532,7 +539,7 @@ public void testDatasourceMetadata() throws Exception expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)) + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)) ) ).anyTimes(); expect(taskQueue.add(capture(captured))).andReturn(true); @@ -543,11 +550,11 @@ public void testDatasourceMetadata() throws Exception verifyAll(); KafkaIndexTask task = captured.getValue(); - KafkaIOConfig taskConfig = task.getIOConfig(); + KafkaIndexTaskIOConfig taskConfig = task.getIOConfig(); Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); - Assert.assertEquals(10L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(20L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(30L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(10L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(20L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(30L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); } @Test(expected = ISE.class) @@ -560,7 +567,7 @@ public void testBadMetadataOffsets() throws Exception expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)) + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)) ) ).anyTimes(); replayAll(); @@ -580,8 +587,8 @@ public void testKillIncompatibleTasks() throws Exception "id1", DATASOURCE, 1, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, 10L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L)), null, null ); @@ -591,8 +598,8 @@ public void testKillIncompatibleTasks() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, 333L, 1, 333L, 2, 333L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 333L, 1, 333L, 2, 333L)), null, null ); @@ -602,8 +609,8 @@ public void testKillIncompatibleTasks() throws Exception "id3", DATASOURCE, 1, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 1L)), - new KafkaPartitions("topic", ImmutableMap.of(0, 333L, 1, 333L, 2, 330L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 1L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 333L, 1, 333L, 2, 330L)), null, null ); @@ -613,8 +620,8 @@ public void testKillIncompatibleTasks() throws Exception "id4", "other-datasource", 2, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, 10L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L)), null, null ); @@ -643,9 +650,12 @@ public void testKillIncompatibleTasks() throws Exception expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); - expect(taskClient.getStatusAsync(anyString())).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)) - .anyTimes(); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -660,7 +670,7 @@ public void testKillIncompatibleTasks() throws Exception TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(2); @@ -681,8 +691,8 @@ public void testKillBadPartitionAssignment() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 2, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -690,8 +700,8 @@ public void testKillBadPartitionAssignment() throws Exception "id2", DATASOURCE, 1, - new KafkaPartitions("topic", ImmutableMap.of(1, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(1, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(1, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(1, Long.MAX_VALUE)), null, null ); @@ -699,8 +709,11 @@ public void testKillBadPartitionAssignment() throws Exception "id3", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -708,8 +721,8 @@ public void testKillBadPartitionAssignment() throws Exception "id4", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE)), null, null ); @@ -717,8 +730,8 @@ public void testKillBadPartitionAssignment() throws Exception "id5", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(1, 0L, 2, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -739,9 +752,12 @@ public void testKillBadPartitionAssignment() throws Exception expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); expect(taskStorage.getTask("id4")).andReturn(Optional.of(id4)).anyTimes(); expect(taskStorage.getTask("id5")).andReturn(Optional.of(id5)).anyTimes(); - expect(taskClient.getStatusAsync(anyString())).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)) - .anyTimes(); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -755,10 +771,10 @@ public void testKillBadPartitionAssignment() throws Exception checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(1); @@ -783,9 +799,12 @@ public void testRequeueTaskWhenFailed() throws Exception expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); - expect(taskClient.getStatusAsync(anyString())).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)) - .anyTimes(); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -797,10 +816,10 @@ public void testRequeueTaskWhenFailed() throws Exception checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .anyTimes(); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .anyTimes(); @@ -815,10 +834,12 @@ public void testRequeueTaskWhenFailed() throws Exception // test that running the main loop again checks the status of the tasks that were created and does nothing if they // are all still running - reset(taskStorage); + EasyMock.reset(taskStorage); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } replay(taskStorage); @@ -830,14 +851,17 @@ public void testRequeueTaskWhenFailed() throws Exception Capture aNewTaskCapture = Capture.newInstance(); List imStillAlive = tasks.subList(0, 3); KafkaIndexTask iHaveFailed = (KafkaIndexTask) tasks.get(3); - reset(taskStorage); - reset(taskQueue); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); expect(taskStorage.getActiveTasks()).andReturn(imStillAlive).anyTimes(); for (Task task : imStillAlive) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } - expect(taskStorage.getStatus(iHaveFailed.getId())).andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); + expect(taskStorage.getStatus(iHaveFailed.getId())) + .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); expect(taskStorage.getTask(iHaveFailed.getId())).andReturn(Optional.of((Task) iHaveFailed)).anyTimes(); expect(taskQueue.add(capture(aNewTaskCapture))).andReturn(true); replay(taskStorage); @@ -865,8 +889,8 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 2, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), now, maxi ); @@ -880,7 +904,7 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception expect(taskStorage.getActiveTasks()).andReturn(existingTasks).anyTimes(); expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(now)).anyTimes(); expect(taskQueue.add(capture(captured))).andReturn(true); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( @@ -891,7 +915,7 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 0L, 2, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(2); @@ -909,24 +933,27 @@ public void testRequeueAdoptedTaskWhenFailed() throws Exception String runningTaskId = captured.getValue().getId(); Capture aNewTaskCapture = Capture.newInstance(); KafkaIndexTask iHaveFailed = (KafkaIndexTask) existingTasks.get(0); - reset(taskStorage); - reset(taskQueue); - reset(taskClient); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); + EasyMock.reset(taskClient); // for the newly created replica task - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(captured.getValue())).anyTimes(); - expect(taskStorage.getStatus(iHaveFailed.getId())).andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); - expect(taskStorage.getStatus(runningTaskId)).andReturn(Optional.of(TaskStatus.running(runningTaskId))).anyTimes(); + expect(taskStorage.getStatus(iHaveFailed.getId())) + .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); + expect(taskStorage.getStatus(runningTaskId)) + .andReturn(Optional.of(TaskStatus.running(runningTaskId))) + .anyTimes(); expect(taskStorage.getTask(iHaveFailed.getId())).andReturn(Optional.of((Task) iHaveFailed)).anyTimes(); expect(taskStorage.getTask(runningTaskId)).andReturn(Optional.of(captured.getValue())).anyTimes(); - expect(taskClient.getStatusAsync(runningTaskId)).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync(runningTaskId)).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync(runningTaskId)).andReturn(Futures.immediateFuture(now)).anyTimes(); expect(taskQueue.add(capture(aNewTaskCapture))).andReturn(true); replay(taskStorage); @@ -962,9 +989,12 @@ public void testQueueNextTasksOnSuccess() throws Exception expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); - expect(taskClient.getStatusAsync(anyString())).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)) - .anyTimes(); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( new KafkaDataSourceMetadata( null @@ -981,27 +1011,32 @@ public void testQueueNextTasksOnSuccess() throws Exception List tasks = captured.getValues(); - reset(taskStorage); - reset(taskClient); + EasyMock.reset(taskStorage); + EasyMock.reset(taskClient); - expect(taskClient.getStatusAsync(anyString())).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)) - .anyTimes(); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); TreeMap> checkpoints1 = new TreeMap<>(); checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); // there would be 4 tasks, 2 for each task group - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } replay(taskStorage); @@ -1016,18 +1051,22 @@ public void testQueueNextTasksOnSuccess() throws Exception Capture shutdownTaskIdCapture = Capture.newInstance(); List imStillRunning = tasks.subList(1, 4); KafkaIndexTask iAmSuccess = (KafkaIndexTask) tasks.get(0); - reset(taskStorage); - reset(taskQueue); - reset(taskClient); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); + EasyMock.reset(taskClient); expect(taskStorage.getActiveTasks()).andReturn(imStillRunning).anyTimes(); for (Task task : imStillRunning) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } - expect(taskStorage.getStatus(iAmSuccess.getId())).andReturn(Optional.of(TaskStatus.success(iAmSuccess.getId()))); + expect(taskStorage.getStatus(iAmSuccess.getId())) + .andReturn(Optional.of(TaskStatus.success(iAmSuccess.getId()))); expect(taskStorage.getTask(iAmSuccess.getId())).andReturn(Optional.of((Task) iAmSuccess)).anyTimes(); expect(taskQueue.add(capture(newTasksCapture))).andReturn(true).times(2); - expect(taskClient.stopAsync(capture(shutdownTaskIdCapture), eq(false))).andReturn(Futures.immediateFuture(true)); + expect(taskClient.stopAsync(capture(shutdownTaskIdCapture), EasyMock.eq(false))) + .andReturn(Futures.immediateFuture(true)); replay(taskStorage); replay(taskQueue); replay(taskClient); @@ -1071,16 +1110,18 @@ public void testBeginPublishAndQueueNextTasks() throws Exception workItems.add(new TestTaskRunnerWorkItem(task, null, location)); } - reset(taskStorage, taskRunner, taskClient, taskQueue); + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); captured = Capture.newInstance(CaptureType.ALL); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - expect(taskClient.getStatusAsync(anyString())) - .andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)) + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) .anyTimes(); expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) @@ -1104,10 +1145,10 @@ public void testBeginPublishAndQueueNextTasks() throws Exception checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); @@ -1119,16 +1160,16 @@ public void testBeginPublishAndQueueNextTasks() throws Exception for (Task task : captured.getValues()) { KafkaIndexTask kafkaIndexTask = (KafkaIndexTask) task; Assert.assertEquals(dataSchema, kafkaIndexTask.getDataSchema()); - Assert.assertEquals(KafkaTuningConfig.copyOf(tuningConfig), kafkaIndexTask.getTuningConfig()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), kafkaIndexTask.getTuningConfig()); - KafkaIOConfig taskConfig = kafkaIndexTask.getIOConfig(); + KafkaIndexTaskIOConfig taskConfig = kafkaIndexTask.getIOConfig(); Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); Assert.assertTrue("isUseTransaction", taskConfig.isUseTransaction()); - Assert.assertEquals(topic, taskConfig.getStartPartitions().getTopic()); - Assert.assertEquals(10L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(20L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(35L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, taskConfig.getStartPartitions().getStream()); + Assert.assertEquals(10L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(20L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(35L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); } } @@ -1144,8 +1185,11 @@ public void testDiscoverExistingPublishingTask() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1165,7 +1209,7 @@ public void testDiscoverExistingPublishingTask() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.PUBLISHING)); expect(taskClient.getCurrentOffsetsAsync("id1", false)) .andReturn(Futures.immediateFuture((Map) ImmutableMap.of(0, 10L, 1, 20L, 2, 30L))); expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); @@ -1173,7 +1217,9 @@ public void testDiscoverExistingPublishingTask() throws Exception TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)); - expect(taskClient.getCheckpoints(anyString(), anyBoolean())).andReturn(checkpoints).anyTimes(); + expect(taskClient.getCheckpoints(EasyMock.anyString(), EasyMock.anyBoolean())) + .andReturn(checkpoints) + .anyTimes(); taskRunner.registerListener(anyObject(TaskRunnerListener.class), anyObject(Executor.class)); replayAll(); @@ -1192,7 +1238,7 @@ public void testDiscoverExistingPublishingTask() throws Exception Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); Assert.assertEquals(NUM_PARTITIONS, (int) payload.getPartitions()); Assert.assertEquals(1, (int) payload.getReplicas()); - Assert.assertEquals(topic, payload.getTopic()); + Assert.assertEquals(topic, payload.getStream()); Assert.assertEquals(0, payload.getActiveTasks().size()); Assert.assertEquals(1, payload.getPublishingTasks().size()); @@ -1204,24 +1250,33 @@ public void testDiscoverExistingPublishingTask() throws Exception KafkaIndexTask capturedTask = captured.getValue(); Assert.assertEquals(dataSchema, capturedTask.getDataSchema()); - Assert.assertEquals(KafkaTuningConfig.copyOf(tuningConfig), capturedTask.getTuningConfig()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), capturedTask.getTuningConfig()); - KafkaIOConfig capturedTaskConfig = capturedTask.getIOConfig(); + KafkaIndexTaskIOConfig capturedTaskConfig = capturedTask.getIOConfig(); Assert.assertEquals(kafkaHost, capturedTaskConfig.getConsumerProperties().get("bootstrap.servers")); Assert.assertEquals("myCustomValue", capturedTaskConfig.getConsumerProperties().get("myCustomKey")); Assert.assertEquals("sequenceName-0", capturedTaskConfig.getBaseSequenceName()); Assert.assertTrue("isUseTransaction", capturedTaskConfig.isUseTransaction()); // check that the new task was created with starting offsets matching where the publishing task finished - Assert.assertEquals(topic, capturedTaskConfig.getStartPartitions().getTopic()); - Assert.assertEquals(10L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(20L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(30L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); - - Assert.assertEquals(topic, capturedTaskConfig.getEndPartitions().getTopic()); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, capturedTaskConfig.getStartPartitions().getStream()); + Assert.assertEquals(10L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(20L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(30L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); + + Assert.assertEquals(topic, capturedTaskConfig.getEndPartitions().getStream()); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(0) + ); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(1) + ); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(2) + ); } @Test @@ -1236,8 +1291,8 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 2, 0L)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -1257,7 +1312,7 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.PUBLISHING)); expect(taskClient.getCurrentOffsetsAsync("id1", false)) .andReturn(Futures.immediateFuture((Map) ImmutableMap.of(0, 10L, 2, 30L))); expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(0, 10L, 2, 30L)); @@ -1280,7 +1335,7 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); Assert.assertEquals(NUM_PARTITIONS, (int) payload.getPartitions()); Assert.assertEquals(1, (int) payload.getReplicas()); - Assert.assertEquals(topic, payload.getTopic()); + Assert.assertEquals(topic, payload.getStream()); Assert.assertEquals(0, payload.getActiveTasks().size()); Assert.assertEquals(1, payload.getPublishingTasks().size()); @@ -1292,24 +1347,33 @@ public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() KafkaIndexTask capturedTask = captured.getValue(); Assert.assertEquals(dataSchema, capturedTask.getDataSchema()); - Assert.assertEquals(KafkaTuningConfig.copyOf(tuningConfig), capturedTask.getTuningConfig()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), capturedTask.getTuningConfig()); - KafkaIOConfig capturedTaskConfig = capturedTask.getIOConfig(); + KafkaIndexTaskIOConfig capturedTaskConfig = capturedTask.getIOConfig(); Assert.assertEquals(kafkaHost, capturedTaskConfig.getConsumerProperties().get("bootstrap.servers")); Assert.assertEquals("myCustomValue", capturedTaskConfig.getConsumerProperties().get("myCustomKey")); Assert.assertEquals("sequenceName-0", capturedTaskConfig.getBaseSequenceName()); Assert.assertTrue("isUseTransaction", capturedTaskConfig.isUseTransaction()); // check that the new task was created with starting offsets matching where the publishing task finished - Assert.assertEquals(topic, capturedTaskConfig.getStartPartitions().getTopic()); - Assert.assertEquals(10L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(30L, (long) capturedTaskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); - - Assert.assertEquals(topic, capturedTaskConfig.getEndPartitions().getTopic()); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(Long.MAX_VALUE, (long) capturedTaskConfig.getEndPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, capturedTaskConfig.getStartPartitions().getStream()); + Assert.assertEquals(10L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(30L, (long) capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); + + Assert.assertEquals(topic, capturedTaskConfig.getEndPartitions().getStream()); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(0) + ); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(1) + ); + Assert.assertEquals( + Long.MAX_VALUE, + (long) capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(2) + ); } @Test @@ -1326,8 +1390,11 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1336,8 +1403,11 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 1L, 1, 2L, 2, 3L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 1L, 1, 2L, 2, 3L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1359,8 +1429,8 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getCurrentOffsetsAsync("id1", false)) .andReturn(Futures.immediateFuture((Map) ImmutableMap.of(0, 1L, 1, 2L, 2, 3L))); @@ -1373,7 +1443,7 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception // since id1 is publishing, so getCheckpoints wouldn't be called for it TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 1L, 1, 2L, 2, 3L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -1393,7 +1463,7 @@ public void testDiscoverExistingPublishingAndReadingTask() throws Exception Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); Assert.assertEquals(NUM_PARTITIONS, (int) payload.getPartitions()); Assert.assertEquals(1, (int) payload.getReplicas()); - Assert.assertEquals(topic, payload.getTopic()); + Assert.assertEquals(topic, payload.getStream()); Assert.assertEquals(1, payload.getActiveTasks().size()); Assert.assertEquals(1, payload.getPublishingTasks().size()); @@ -1443,25 +1513,27 @@ public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception List tasks = captured.getValues(); - reset(taskStorage, taskClient, taskQueue); + EasyMock.reset(taskStorage, taskClient, taskQueue); TreeMap> checkpoints1 = new TreeMap<>(); checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); expect(taskClient.getStatusAsync(task.getId())) - .andReturn(Futures.immediateFuture(KafkaIndexTask.Status.NOT_STARTED)); + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)); expect(taskClient.getStartTimeAsync(task.getId())) .andReturn(Futures.immediateFailedFuture(new RuntimeException())); taskQueue.shutdown(task.getId(), "Task [%s] failed to return start time, killing task", task.getId()); @@ -1504,28 +1576,30 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception workItems.add(new TestTaskRunnerWorkItem(task, null, location)); } - reset(taskStorage, taskRunner, taskClient, taskQueue); + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); TreeMap> checkpoints1 = new TreeMap<>(); checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); captured = Capture.newInstance(CaptureType.ALL); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - expect(taskClient.getStatusAsync(anyString())) - .andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)) + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) .anyTimes(); expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) @@ -1550,9 +1624,9 @@ public void testKillUnresponsiveTasksWhilePausing() throws Exception verifyAll(); for (Task task : captured.getValues()) { - KafkaIOConfig taskConfig = ((KafkaIndexTask) task).getIOConfig(); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + KafkaIndexTaskIOConfig taskConfig = ((KafkaIndexTask) task).getIOConfig(); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); } } @@ -1588,28 +1662,30 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception workItems.add(new TestTaskRunnerWorkItem(task, null, location)); } - reset(taskStorage, taskRunner, taskClient, taskQueue); + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); TreeMap> checkpoints1 = new TreeMap<>(); checkpoints1.put(0, ImmutableMap.of(0, 0L, 2, 0L)); TreeMap> checkpoints2 = new TreeMap<>(); checkpoints2.put(0, ImmutableMap.of(1, 0L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints1)) .times(2); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints2)) .times(2); captured = Capture.newInstance(CaptureType.ALL); expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); for (Task task : tasks) { - expect(taskStorage.getStatus(task.getId())).andReturn(Optional.of(TaskStatus.running(task.getId()))).anyTimes(); + expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); } expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); - expect(taskClient.getStatusAsync(anyString())) - .andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)) + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) .anyTimes(); expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) @@ -1641,9 +1717,9 @@ public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception verifyAll(); for (Task task : captured.getValues()) { - KafkaIOConfig taskConfig = ((KafkaIndexTask) task).getIOConfig(); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + KafkaIndexTaskIOConfig taskConfig = ((KafkaIndexTask) task).getIOConfig(); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); } } @@ -1683,8 +1759,11 @@ public void testStopGracefully() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1693,8 +1772,11 @@ public void testStopGracefully() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1703,8 +1785,11 @@ public void testStopGracefully() throws Exception "id3", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1728,9 +1813,9 @@ public void testStopGracefully() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); @@ -1738,10 +1823,10 @@ public void testStopGracefully() throws Exception // getCheckpoints will not be called for id1 as it is in publishing state TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -1752,7 +1837,7 @@ public void testStopGracefully() throws Exception supervisor.runInternal(); verifyAll(); - reset(taskRunner, taskClient, taskQueue); + EasyMock.reset(taskRunner, taskClient, taskQueue); expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); expect(taskClient.pauseAsync("id2")) .andReturn(Futures.immediateFuture((Map) ImmutableMap.of(0, 15L, 1, 25L, 2, 30L))); @@ -1784,7 +1869,7 @@ public void testResetNoTasks() throws Exception supervisor.runInternal(); verifyAll(); - reset(indexerMetadataStorageCoordinator); + EasyMock.reset(indexerMetadataStorageCoordinator); expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); replay(indexerMetadataStorageCoordinator); @@ -1811,26 +1896,27 @@ public void testResetDataSourceMetadata() throws Exception Capture captureDataSource = EasyMock.newCapture(); Capture captureDataSourceMetadata = EasyMock.newCapture(); - KafkaDataSourceMetadata kafkaDataSourceMetadata = new KafkaDataSourceMetadata(new KafkaPartitions( + KafkaDataSourceMetadata kafkaDataSourceMetadata = new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( topic, ImmutableMap.of(0, 1000L, 1, 1000L, 2, 1000L) )); - KafkaDataSourceMetadata resetMetadata = new KafkaDataSourceMetadata(new KafkaPartitions( + KafkaDataSourceMetadata resetMetadata = new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( topic, ImmutableMap.of(1, 1000L, 2, 1000L) )); - KafkaDataSourceMetadata expectedMetadata = new KafkaDataSourceMetadata(new KafkaPartitions( + KafkaDataSourceMetadata expectedMetadata = new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( topic, ImmutableMap.of(0, 1000L) )); - reset(indexerMetadataStorageCoordinator); - expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(kafkaDataSourceMetadata); + EasyMock.reset(indexerMetadataStorageCoordinator); + expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)) + .andReturn(kafkaDataSourceMetadata); expect(indexerMetadataStorageCoordinator.resetDataSourceMetadata( - EasyMock.capture(captureDataSource), - EasyMock.capture(captureDataSourceMetadata) + capture(captureDataSource), + capture(captureDataSourceMetadata) )).andReturn(true); replay(indexerMetadataStorageCoordinator); @@ -1838,7 +1924,7 @@ public void testResetDataSourceMetadata() throws Exception supervisor.resetInternal(resetMetadata); } catch (NullPointerException npe) { - // Expected as there will be an attempt to reset partitionGroups offsets to NOT_SET + // Expected as there will be an attempt to EasyMock.reset partitionGroups offsets to NOT_SET // however there would be no entries in the map as we have not put nay data in kafka Assert.assertTrue(npe.getCause() == null); } @@ -1863,12 +1949,12 @@ public void testResetNoDataSourceMetadata() throws Exception supervisor.runInternal(); verifyAll(); - KafkaDataSourceMetadata resetMetadata = new KafkaDataSourceMetadata(new KafkaPartitions( + KafkaDataSourceMetadata resetMetadata = new KafkaDataSourceMetadata(new SeekableStreamPartitions<>( topic, ImmutableMap.of(1, 1000L, 2, 1000L) )); - reset(indexerMetadataStorageCoordinator); + EasyMock.reset(indexerMetadataStorageCoordinator); // no DataSourceMetadata in metadata store expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(null); replay(indexerMetadataStorageCoordinator); @@ -1891,8 +1977,11 @@ public void testResetRunningTasks() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1901,8 +1990,11 @@ public void testResetRunningTasks() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1911,8 +2003,11 @@ public void testResetRunningTasks() throws Exception "id3", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1936,19 +2031,19 @@ public void testResetRunningTasks() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -1959,7 +2054,7 @@ public void testResetRunningTasks() throws Exception supervisor.runInternal(); verifyAll(); - reset(taskQueue, indexerMetadataStorageCoordinator); + EasyMock.reset(taskQueue, indexerMetadataStorageCoordinator); expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); taskQueue.shutdown("id2", "DataSourceMetadata is not found while reset"); taskQueue.shutdown("id3", "DataSourceMetadata is not found while reset"); @@ -1979,8 +2074,11 @@ public void testNoDataIngestionTasks() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1989,8 +2087,11 @@ public void testNoDataIngestionTasks() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -1999,8 +2100,11 @@ public void testNoDataIngestionTasks() throws Exception "id3", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -2019,22 +2123,22 @@ public void testNoDataIngestionTasks() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -2045,7 +2149,7 @@ public void testNoDataIngestionTasks() throws Exception supervisor.runInternal(); verifyAll(); - reset(taskQueue, indexerMetadataStorageCoordinator); + EasyMock.reset(taskQueue, indexerMetadataStorageCoordinator); expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); taskQueue.shutdown("id1", "DataSourceMetadata is not found while reset"); taskQueue.shutdown("id2", "DataSourceMetadata is not found while reset"); @@ -2066,8 +2170,8 @@ public void testCheckpointForInactiveTaskGroup() "id1", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2076,8 +2180,8 @@ public void testCheckpointForInactiveTaskGroup() "id2", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2086,8 +2190,8 @@ public void testCheckpointForInactiveTaskGroup() "id3", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2109,12 +2213,12 @@ public void testCheckpointForInactiveTaskGroup() expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); - expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)) - .andReturn(new KafkaDataSourceMetadata(null)) - .anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect( + indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(new KafkaDataSourceMetadata(null) + ).anyTimes(); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(Status.READING)); final DateTime startTime = DateTimes.nowUtc(); expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(startTime)); @@ -2123,13 +2227,13 @@ public void testCheckpointForInactiveTaskGroup() final TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -2144,8 +2248,8 @@ public void testCheckpointForInactiveTaskGroup() supervisor.checkpoint( 0, ((KafkaIndexTask) id1).getIOConfig().getBaseSequenceName(), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, checkpoints.get(0))), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, fakeCheckpoints)) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, checkpoints.get(0))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, fakeCheckpoints)) ); while (supervisor.getNoticesQueueSize() > 0) { @@ -2160,7 +2264,8 @@ public void testCheckpointForInactiveTaskGroup() } @Test(timeout = 60_000L) - public void testCheckpointForUnknownTaskGroup() throws InterruptedException + public void testCheckpointForUnknownTaskGroup() + throws InterruptedException { supervisor = getSupervisor(2, 1, true, "PT1S", null, null, false); //not adding any events @@ -2168,8 +2273,8 @@ public void testCheckpointForUnknownTaskGroup() throws InterruptedException "id1", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2178,8 +2283,8 @@ public void testCheckpointForUnknownTaskGroup() throws InterruptedException "id2", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2188,8 +2293,8 @@ public void testCheckpointForUnknownTaskGroup() throws InterruptedException "id3", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), null, null ); @@ -2214,8 +2319,8 @@ public void testCheckpointForUnknownTaskGroup() throws InterruptedException supervisor.checkpoint( 0, ((KafkaIndexTask) id1).getIOConfig().getBaseSequenceName(), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, Collections.emptyMap())), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, Collections.emptyMap())) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, Collections.emptyMap())), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, Collections.emptyMap())) ); while (supervisor.getNoticesQueueSize() > 0) { @@ -2231,7 +2336,7 @@ public void testCheckpointForUnknownTaskGroup() throws InterruptedException Assert.assertTrue(serviceEmitter.getStackTrace() .startsWith("org.apache.druid.java.util.common.ISE: WTH?! cannot find")); Assert.assertEquals( - "WTH?! cannot find taskGroup [0] among all taskGroups [{}]", + "WTH?! cannot find taskGroup [0] among all activelyReadingTaskGroups [{}]", serviceEmitter.getExceptionMessage() ); Assert.assertEquals(ISE.class, serviceEmitter.getExceptionClass()); @@ -2247,8 +2352,8 @@ public void testCheckpointWithNullTaskGroupId() "id1", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE)), null, null ); @@ -2257,8 +2362,8 @@ public void testCheckpointWithNullTaskGroupId() "id2", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE)), null, null ); @@ -2267,8 +2372,8 @@ public void testCheckpointWithNullTaskGroupId() "id3", DATASOURCE, 0, - new KafkaPartitions(topic, ImmutableMap.of(0, 0L)), - new KafkaPartitions(topic, ImmutableMap.of(0, Long.MAX_VALUE)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, 0L)), + new SeekableStreamPartitions<>(topic, ImmutableMap.of(0, Long.MAX_VALUE)), null, null ); @@ -2286,19 +2391,25 @@ public void testCheckpointWithNullTaskGroupId() indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(new KafkaDataSourceMetadata(null) ).anyTimes(); taskRunner.registerListener(anyObject(TaskRunnerListener.class), anyObject(Executor.class)); - expect(taskClient.getStatusAsync(anyString())) - .andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)) + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) .anyTimes(); final TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 0L)); - expect(taskClient.getCheckpointsAsync(anyString(), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.anyString(), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(3); - expect(taskClient.getStartTimeAsync(anyString())).andReturn(Futures.immediateFuture(DateTimes.nowUtc())).anyTimes(); - expect(taskClient.pauseAsync(anyString())) + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + expect(taskClient.pauseAsync(EasyMock.anyString())) .andReturn(Futures.immediateFuture(ImmutableMap.of(0, 10L))) .anyTimes(); - expect(taskClient.setEndOffsetsAsync(anyString(), EasyMock.eq(ImmutableMap.of(0, 10L)), anyBoolean())) + expect(taskClient.setEndOffsetsAsync( + EasyMock.anyString(), + EasyMock.eq(ImmutableMap.of(0, 10L)), + EasyMock.anyBoolean() + )) .andReturn(Futures.immediateFuture(true)) .anyTimes(); @@ -2313,8 +2424,8 @@ public void testCheckpointWithNullTaskGroupId() supervisor.checkpoint( null, ((KafkaIndexTask) id1).getIOConfig().getBaseSequenceName(), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, checkpoints.get(0))), - new KafkaDataSourceMetadata(new KafkaPartitions(topic, newCheckpoints.get(0))) + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, checkpoints.get(0))), + new KafkaDataSourceMetadata(new SeekableStreamPartitions<>(topic, newCheckpoints.get(0))) ); while (supervisor.getNoticesQueueSize() > 0) { @@ -2367,8 +2478,11 @@ public void testSuspendedRunningTasks() throws Exception "id1", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 0L, 1, 0L, 2, 0L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -2377,8 +2491,11 @@ public void testSuspendedRunningTasks() throws Exception "id2", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -2387,8 +2504,11 @@ public void testSuspendedRunningTasks() throws Exception "id3", DATASOURCE, 0, - new KafkaPartitions("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), - new KafkaPartitions("topic", ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE)), + new SeekableStreamPartitions<>("topic", ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)), + new SeekableStreamPartitions<>( + "topic", + ImmutableMap.of(0, Long.MAX_VALUE, 1, Long.MAX_VALUE, 2, Long.MAX_VALUE) + ), null, null ); @@ -2412,9 +2532,12 @@ public void testSuspendedRunningTasks() throws Exception null ) ).anyTimes(); - expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.PUBLISHING)); - expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); - expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(KafkaIndexTask.Status.READING)); + expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")) + .andReturn(Futures.immediateFuture(Status.READING)); expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); @@ -2422,10 +2545,10 @@ public void testSuspendedRunningTasks() throws Exception // getCheckpoints will not be called for id1 as it is in publishing state TreeMap> checkpoints = new TreeMap<>(); checkpoints.put(0, ImmutableMap.of(0, 10L, 1, 20L, 2, 30L)); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); - expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), anyBoolean())) + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) .andReturn(Futures.immediateFuture(checkpoints)) .times(1); @@ -2461,7 +2584,7 @@ public void testResetSuspended() throws Exception supervisor.runInternal(); verifyAll(); - reset(indexerMetadataStorageCoordinator); + EasyMock.reset(indexerMetadataStorageCoordinator); expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); replay(indexerMetadataStorageCoordinator); @@ -2538,9 +2661,9 @@ public void testFailedInitializationAndRecovery() throws Exception KafkaIndexTask task = captured.getValue(); Assert.assertEquals(dataSchema, task.getDataSchema()); - Assert.assertEquals(KafkaTuningConfig.copyOf(tuningConfig), task.getTuningConfig()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), task.getTuningConfig()); - KafkaIOConfig taskConfig = task.getIOConfig(); + KafkaIndexTaskIOConfig taskConfig = task.getIOConfig(); Assert.assertEquals(kafkaHost, taskConfig.getConsumerProperties().get("bootstrap.servers")); Assert.assertEquals("myCustomValue", taskConfig.getConsumerProperties().get("myCustomKey")); Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); @@ -2549,19 +2672,19 @@ public void testFailedInitializationAndRecovery() throws Exception Assert.assertFalse("maximumMessageTime", taskConfig.getMaximumMessageTime().isPresent()); Assert.assertFalse("skipOffsetGaps", taskConfig.isSkipOffsetGaps()); - Assert.assertEquals(topic, taskConfig.getStartPartitions().getTopic()); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, taskConfig.getStartPartitions().getStream()); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(0L, (long) taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(2)); - Assert.assertEquals(topic, taskConfig.getEndPartitions().getTopic()); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(0)); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(1)); - Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionOffsetMap().get(2)); + Assert.assertEquals(topic, taskConfig.getEndPartitions().getStream()); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(0)); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(1)); + Assert.assertEquals(Long.MAX_VALUE, (long) taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(2)); } @Test - public void testGetCurrentTotalStats() throws Exception + public void testGetCurrentTotalStats() { supervisor = getSupervisor(1, 2, true, "PT1H", null, null, false); supervisor.addTaskGroupToActivelyReadingTaskGroup( @@ -2569,7 +2692,8 @@ public void testGetCurrentTotalStats() throws Exception ImmutableMap.of(0, 0L), Optional.absent(), Optional.absent(), - ImmutableSet.of("task1") + ImmutableSet.of("task1"), + ImmutableSet.of() ); supervisor.addTaskGroupToPendingCompletionTaskGroup( @@ -2577,7 +2701,8 @@ public void testGetCurrentTotalStats() throws Exception ImmutableMap.of(0, 0L), Optional.absent(), Optional.absent(), - ImmutableSet.of("task2") + ImmutableSet.of("task2"), + ImmutableSet.of() ); expect(taskClient.getMovingAveragesAsync("task1")).andReturn(Futures.immediateFuture(ImmutableMap.of( @@ -2762,8 +2887,8 @@ private KafkaIndexTask createKafkaIndexTask( String id, String dataSource, int taskGroupId, - KafkaPartitions startPartitions, - KafkaPartitions endPartitions, + SeekableStreamPartitions startPartitions, + SeekableStreamPartitions endPartitions, DateTime minimumMessageTime, DateTime maximumMessageTime ) @@ -2773,12 +2898,12 @@ private KafkaIndexTask createKafkaIndexTask( null, getDataSchema(dataSource), tuningConfig, - new KafkaIOConfig( + new KafkaIndexTaskIOConfig( taskGroupId, "sequenceName-" + taskGroupId, startPartitions, endPartitions, - ImmutableMap.of(), + ImmutableMap.of("bootstrap.servers", kafkaHost), true, minimumMessageTime, maximumMessageTime, @@ -2860,4 +2985,6 @@ protected String generateSequenceName( return StringUtils.format("sequenceName-%d", groupId); } } + + } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfigTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfigTest.java index 0a3ea95fea54..45470ffe10de 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfigTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTuningConfigTest.java @@ -107,7 +107,7 @@ public void testSerdeWithNonDefaults() throws Exception Assert.assertEquals(100, config.getMaxRowsInMemory()); Assert.assertEquals(100, config.getMaxRowsPerSegment()); Assert.assertEquals(new Period("PT1H"), config.getIntermediatePersistPeriod()); - Assert.assertEquals(0, config.getMaxPendingPersists()); + Assert.assertEquals(100, config.getMaxPendingPersists()); Assert.assertEquals(true, config.isReportParseExceptions()); Assert.assertEquals(100, config.getHandoffConditionTimeout()); Assert.assertEquals(12, (int) config.getWorkerThreads()); @@ -117,4 +117,5 @@ public void testSerdeWithNonDefaults() throws Exception Assert.assertEquals(Duration.standardSeconds(95), config.getShutdownTimeout()); Assert.assertEquals(Duration.standardSeconds(20), config.getOffsetFetchPeriod()); } + } diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/test/TestBroker.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/test/TestBroker.java index a57d22b1446c..6659f92ebbf0 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/test/TestBroker.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/test/TestBroker.java @@ -49,7 +49,7 @@ public class TestBroker implements Closeable private final File directory; private final boolean directoryCleanup; private final int id; - private final Map brokerProps; + private final Map brokerProps; private volatile KafkaServer server; @@ -57,7 +57,7 @@ public TestBroker( String zookeeperConnect, @Nullable File directory, int id, - Map brokerProps + Map brokerProps ) { this.zookeeperConnect = zookeeperConnect; @@ -81,7 +81,12 @@ public void start() final KafkaConfig config = new KafkaConfig(props); - server = new KafkaServer(config, Time.SYSTEM, Some.apply(StringUtils.format("TestingBroker[%d]-", id)), List$.MODULE$.empty()); + server = new KafkaServer( + config, + Time.SYSTEM, + Some.apply(StringUtils.format("TestingBroker[%d]-", id)), + List$.MODULE$.empty() + ); server.startup(); } diff --git a/extensions-core/kinesis-indexing-service/pom.xml b/extensions-core/kinesis-indexing-service/pom.xml new file mode 100644 index 000000000000..db79d856671c --- /dev/null +++ b/extensions-core/kinesis-indexing-service/pom.xml @@ -0,0 +1,109 @@ + + + + + 4.0.0 + + org.apache.druid.extensions + druid-kinesis-indexing-service + druid-kinesis-indexing-service + druid-kinesis-indexing-service + + + org.apache.druid + druid + 0.13.0-incubating-SNAPSHOT + ../../pom.xml + + + + + org.apache.druid + druid-core + ${project.parent.version} + provided + + + org.apache.druid + druid-indexing-service + ${project.parent.version} + provided + + + org.apache.druid + druid-server + ${project.parent.version} + provided + + + com.amazonaws + aws-java-sdk-kinesis + ${aws.sdk.version} + + + + com.amazonaws + aws-java-sdk-sts + ${aws.sdk.version} + + + + com.amazonaws + amazon-kinesis-client + 1.9.2 + provided + + + + + org.easymock + easymock + test + + + junit + junit + test + + + org.apache.druid + druid-server + ${project.parent.version} + test-jar + test + + + org.apache.druid + druid-processing + ${project.parent.version} + test-jar + test + + + org.apache.druid + druid-indexing-service + ${project.parent.version} + test-jar + test + + + + diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadata.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadata.java new file mode 100644 index 000000000000..4592c1b67740 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadata.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; + +import java.util.Map; + +public class KinesisDataSourceMetadata extends SeekableStreamDataSourceMetadata +{ + @JsonCreator + public KinesisDataSourceMetadata( + @JsonProperty("partitions") SeekableStreamPartitions kinesisPartitions + ) + { + super(kinesisPartitions); + } + + @Override + protected KinesisDataSourceMetadata createConcreteDataSourceMetaData(String streamName, Map newMap) + { + return new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(streamName, newMap)); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTask.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTask.java new file mode 100644 index 000000000000..3467481762af --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTask.java @@ -0,0 +1,126 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JacksonInject; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.TaskResource; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; +import org.apache.druid.server.security.AuthorizerMapper; + +import java.util.Map; + +public class KinesisIndexTask extends SeekableStreamIndexTask +{ + private static final String TYPE = "index_kinesis"; + + private final AWSCredentialsConfig awsCredentialsConfig; + + @JsonCreator + public KinesisIndexTask( + @JsonProperty("id") String id, + @JsonProperty("resource") TaskResource taskResource, + @JsonProperty("dataSchema") DataSchema dataSchema, + @JsonProperty("tuningConfig") KinesisIndexTaskTuningConfig tuningConfig, + @JsonProperty("ioConfig") KinesisIndexTaskIOConfig ioConfig, + @JsonProperty("context") Map context, + @JacksonInject ChatHandlerProvider chatHandlerProvider, + @JacksonInject AuthorizerMapper authorizerMapper, + @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory, + @JacksonInject AWSCredentialsConfig awsCredentialsConfig + ) + { + super( + id == null ? getFormattedId(dataSchema.getDataSource(), TYPE) : id, + taskResource, + dataSchema, + tuningConfig, + ioConfig, + context, + chatHandlerProvider, + authorizerMapper, + rowIngestionMetersFactory, + getFormattedGroupId(dataSchema.getDataSource(), TYPE) + ); + this.awsCredentialsConfig = awsCredentialsConfig; + } + + + @Override + protected SeekableStreamIndexTaskRunner createTaskRunner() + { + return new KinesisIndexTaskRunner( + this, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory + ); + } + + @Override + protected KinesisRecordSupplier newTaskRecordSupplier() + throws RuntimeException + { + KinesisIndexTaskIOConfig ioConfig = ((KinesisIndexTaskIOConfig) super.ioConfig); + KinesisIndexTaskTuningConfig tuningConfig = ((KinesisIndexTaskTuningConfig) super.tuningConfig); + int fetchThreads = tuningConfig.getFetchThreads() != null + ? tuningConfig.getFetchThreads() + : Math.max(1, ioConfig.getStartPartitions().getPartitionSequenceNumberMap().size()); + + return new KinesisRecordSupplier( + KinesisRecordSupplier.getAmazonKinesisClient( + ioConfig.getEndpoint(), + awsCredentialsConfig, + ioConfig.getAwsAssumedRoleArn(), + ioConfig.getAwsExternalId() + ), + ioConfig.getRecordsPerFetch(), + ioConfig.getFetchDelayMillis(), + fetchThreads, + ioConfig.isDeaggregate(), + tuningConfig.getRecordBufferSize(), + tuningConfig.getRecordBufferOfferTimeout(), + tuningConfig.getRecordBufferFullWait(), + tuningConfig.getFetchSequenceNumberTimeout(), + tuningConfig.getMaxRecordsPerPoll() + ); + } + + @Override + @JsonProperty("ioConfig") + public KinesisIndexTaskIOConfig getIOConfig() + { + return (KinesisIndexTaskIOConfig) super.getIOConfig(); + } + + @Override + public String getType() + { + return TYPE; + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClient.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClient.java new file mode 100644 index 000000000000..8f1c66a1acde --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClient.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClient; +import org.apache.druid.java.util.http.client.HttpClient; +import org.joda.time.Duration; + +public class KinesisIndexTaskClient extends SeekableStreamIndexTaskClient +{ + KinesisIndexTaskClient( + HttpClient httpClient, + ObjectMapper jsonMapper, + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ) + { + super( + httpClient, + jsonMapper, + taskInfoProvider, + dataSource, + numThreads, + httpTimeout, + numRetries + ); + } + + @Override + protected Class getPartitionType() + { + return String.class; + } + + @Override + protected Class getSequenceType() + { + return String.class; + } + +} + diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java new file mode 100644 index 000000000000..6d443464da6b --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.Inject; +import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; +import org.apache.druid.java.util.http.client.HttpClient; +import org.joda.time.Duration; + +public class KinesisIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory +{ + @Inject + public KinesisIndexTaskClientFactory( + @EscalatedGlobal HttpClient httpClient, + @Json ObjectMapper mapper + ) + { + super(httpClient, mapper); + } + + @Override + public KinesisIndexTaskClient build( + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ) + { + return new KinesisIndexTaskClient( + getHttpClient(), + getMapper(), + taskInfoProvider, + dataSource, + numThreads, + httpTimeout, + numRetries + ); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskIOConfig.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskIOConfig.java new file mode 100644 index 000000000000..307e971bcfd3 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskIOConfig.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import java.util.Set; + +public class KinesisIndexTaskIOConfig extends SeekableStreamIndexTaskIOConfig +{ + public static final int DEFAULT_RECORDS_PER_FETCH = 4000; + public static final int DEFAULT_FETCH_DELAY_MILLIS = 0; + + private final String endpoint; + private final Integer recordsPerFetch; + private final Integer fetchDelayMillis; + + private final String awsAssumedRoleArn; + private final String awsExternalId; + private final boolean deaggregate; + + @JsonCreator + public KinesisIndexTaskIOConfig( + @JsonProperty("taskGroupId") @Nullable Integer taskGroupId, + @JsonProperty("baseSequenceName") String baseSequenceName, + @JsonProperty("startPartitions") SeekableStreamPartitions startPartitions, + @JsonProperty("endPartitions") SeekableStreamPartitions endPartitions, + @JsonProperty("useTransaction") Boolean useTransaction, + @JsonProperty("minimumMessageTime") DateTime minimumMessageTime, + @JsonProperty("maximumMessageTime") DateTime maximumMessageTime, + @JsonProperty("endpoint") String endpoint, + @JsonProperty("recordsPerFetch") Integer recordsPerFetch, + @JsonProperty("fetchDelayMillis") Integer fetchDelayMillis, + @JsonProperty("exclusiveStartSequenceNumberPartitions") Set exclusiveStartSequenceNumberPartitions, + @JsonProperty("awsAssumedRoleArn") String awsAssumedRoleArn, + @JsonProperty("awsExternalId") String awsExternalId, + @JsonProperty("deaggregate") boolean deaggregate + ) + { + super( + taskGroupId, + baseSequenceName, + startPartitions, + endPartitions, + useTransaction, + minimumMessageTime, + maximumMessageTime, + true, + exclusiveStartSequenceNumberPartitions + ); + Preconditions.checkArgument(endPartitions.getPartitionSequenceNumberMap() + .values() + .stream() + .noneMatch(x -> x.equals(KinesisSequenceNumber.END_OF_SHARD_MARKER))); + + this.endpoint = Preconditions.checkNotNull(endpoint, "endpoint"); + this.recordsPerFetch = recordsPerFetch != null ? recordsPerFetch : DEFAULT_RECORDS_PER_FETCH; + this.fetchDelayMillis = fetchDelayMillis != null ? fetchDelayMillis : DEFAULT_FETCH_DELAY_MILLIS; + this.awsAssumedRoleArn = awsAssumedRoleArn; + this.awsExternalId = awsExternalId; + this.deaggregate = deaggregate; + } + + @JsonProperty + public String getEndpoint() + { + return endpoint; + } + + @JsonProperty + public int getRecordsPerFetch() + { + return recordsPerFetch; + } + + @JsonProperty + public int getFetchDelayMillis() + { + return fetchDelayMillis; + } + + @JsonProperty + public String getAwsAssumedRoleArn() + { + return awsAssumedRoleArn; + } + + @JsonProperty + public String getAwsExternalId() + { + return awsExternalId; + } + + @JsonProperty + public boolean isDeaggregate() + { + return deaggregate; + } + + @Override + public String toString() + { + return "KinesisIndexTaskIOConfig{" + + "baseSequenceName='" + getBaseSequenceName() + '\'' + + ", startPartitions=" + getStartPartitions() + + ", endPartitions=" + getEndPartitions() + + ", useTransaction=" + isUseTransaction() + + ", minimumMessageTime=" + getMinimumMessageTime() + + ", maximumMessageTime=" + getMaximumMessageTime() + + ", endpoint='" + endpoint + '\'' + + ", recordsPerFetch=" + recordsPerFetch + + ", fetchDelayMillis=" + fetchDelayMillis + + ", exclusiveStartSequenceNumberPartitions=" + getExclusiveStartSequenceNumberPartitions() + + ", awsAssumedRoleArn='" + awsAssumedRoleArn + '\'' + + ", awsExternalId='" + awsExternalId + '\'' + + ", deaggregate=" + deaggregate + + '}'; + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskRunner.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskRunner.java new file mode 100644 index 000000000000..3e7e5e7aeddd --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskRunner.java @@ -0,0 +1,199 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Optional; +import org.apache.druid.data.input.impl.InputRowParser; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.utils.CircularBuffer; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.stream.Collectors; + +public class KinesisIndexTaskRunner extends SeekableStreamIndexTaskRunner +{ + private static final EmittingLogger log = new EmittingLogger(KinesisIndexTaskRunner.class); + private static final long POLL_TIMEOUT = 100; + + private final KinesisIndexTask task; + + KinesisIndexTaskRunner( + KinesisIndexTask task, + InputRowParser parser, + AuthorizerMapper authorizerMapper, + Optional chatHandlerProvider, + CircularBuffer savedParseExceptions, + RowIngestionMetersFactory rowIngestionMetersFactory + ) + { + super( + task, + parser, + authorizerMapper, + chatHandlerProvider, + savedParseExceptions, + rowIngestionMetersFactory + ); + this.task = task; + } + + + @Override + protected String getSequenceNumberToStoreAfterRead(String sequenceNumber) + { + return sequenceNumber; + } + + @Nonnull + @Override + protected List> getRecords( + RecordSupplier recordSupplier, TaskToolbox toolbox + ) + { + return recordSupplier.poll(POLL_TIMEOUT); + } + + @Override + protected SeekableStreamPartitions deserializeSeekableStreamPartitionsFromMetadata( + ObjectMapper mapper, + Object object + ) + { + return mapper.convertValue(object, mapper.getTypeFactory().constructParametrizedType( + SeekableStreamPartitions.class, + SeekableStreamPartitions.class, + String.class, + String.class + )); + } + + @Override + protected SeekableStreamDataSourceMetadata createDataSourceMetadata( + SeekableStreamPartitions partitions + ) + { + return new KinesisDataSourceMetadata(partitions); + } + + @Override + protected OrderedSequenceNumber createSequenceNumber(String sequenceNumber) + { + return KinesisSequenceNumber.of(sequenceNumber); + } + + @Override + protected void possiblyResetDataSourceMetadata( + TaskToolbox toolbox, + RecordSupplier recordSupplier, + Set> assignment, + Map currOffsets + ) + { + if (!task.getTuningConfig().isSkipSequenceNumberAvailabilityCheck()) { + for (final StreamPartition streamPartition : assignment) { + String sequence = currOffsets.get(streamPartition.getPartitionId()); + String earliestSequenceNumber = recordSupplier.getEarliestSequenceNumber(streamPartition); + if (earliestSequenceNumber == null + || createSequenceNumber(earliestSequenceNumber).compareTo(createSequenceNumber(sequence)) > 0) { + if (task.getTuningConfig().isResetOffsetAutomatically()) { + log.info("Attempting to reset sequences automatically for all partitions"); + try { + sendResetRequestAndWait( + assignment.stream() + .collect(Collectors.toMap(x -> x, x -> currOffsets.get(x.getPartitionId()))), + toolbox + ); + } + catch (IOException e) { + throw new ISE(e, "Exception while attempting to automatically reset sequences"); + } + } else { + throw new ISE( + "Starting sequenceNumber [%s] is no longer available for partition [%s] (earliest: [%s]) and resetOffsetAutomatically is not enabled", + sequence, + streamPartition.getPartitionId(), + earliestSequenceNumber + ); + } + } + } + } + } + + @Override + protected boolean isEndSequenceOffsetsExclusive() + { + return false; + } + + @Override + protected boolean isStartingSequenceOffsetsExclusive() + { + return true; + } + + @Override + protected boolean isEndOfShard(String seqNum) + { + return KinesisSequenceNumber.END_OF_SHARD_MARKER.equals(seqNum); + } + + @Nullable + @Override + protected TreeMap> getCheckPointsFromContext( + TaskToolbox toolbox, + String checkpointsString + ) throws IOException + { + if (checkpointsString != null) { + log.info("Checkpoints [%s]", checkpointsString); + return toolbox.getObjectMapper().readValue( + checkpointsString, + new TypeReference>>() + { + } + ); + } else { + return null; + } + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfig.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfig.java new file mode 100644 index 000000000000..95ee278985fe --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfig.java @@ -0,0 +1,278 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import com.google.common.base.Preconditions; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory; +import org.joda.time.Period; + +import javax.annotation.Nullable; +import java.io.File; +import java.util.Objects; + +@JsonTypeName("KinesisTuningConfig") +public class KinesisIndexTaskTuningConfig extends SeekableStreamIndexTaskTuningConfig +{ + private static final int DEFAULT_RECORD_BUFFER_SIZE = 10000; + private static final int DEFAULT_RECORD_BUFFER_OFFER_TIMEOUT = 5000; + private static final int DEFAULT_RECORD_BUFFER_FULL_WAIT = 5000; + private static final int DEFAULT_FETCH_SEQUENCE_NUMBER_TIMEOUT = 20000; + private static final int DEFAULT_MAX_RECORDS_PER_POLL = 100; + + private final int recordBufferSize; + private final int recordBufferOfferTimeout; + private final int recordBufferFullWait; + private final int fetchSequenceNumberTimeout; + private final Integer fetchThreads; + private final int maxRecordsPerPoll; + + @JsonCreator + public KinesisIndexTaskTuningConfig( + @JsonProperty("maxRowsInMemory") Integer maxRowsInMemory, + @JsonProperty("maxBytesInMemory") Long maxBytesInMemory, + @JsonProperty("maxRowsPerSegment") Integer maxRowsPerSegment, + @JsonProperty("maxTotalRows") Long maxTotalRows, + @JsonProperty("intermediatePersistPeriod") Period intermediatePersistPeriod, + @JsonProperty("basePersistDirectory") File basePersistDirectory, + @JsonProperty("maxPendingPersists") Integer maxPendingPersists, + @JsonProperty("indexSpec") IndexSpec indexSpec, + @JsonProperty("buildV9Directly") Boolean buildV9Directly, + @JsonProperty("reportParseExceptions") Boolean reportParseExceptions, + @JsonProperty("handoffConditionTimeout") Long handoffConditionTimeout, + @JsonProperty("resetOffsetAutomatically") Boolean resetOffsetAutomatically, + @JsonProperty("skipSequenceNumberAvailabilityCheck") Boolean skipSequenceNumberAvailabilityCheck, + @JsonProperty("recordBufferSize") Integer recordBufferSize, + @JsonProperty("recordBufferOfferTimeout") Integer recordBufferOfferTimeout, + @JsonProperty("recordBufferFullWait") Integer recordBufferFullWait, + @JsonProperty("fetchSequenceNumberTimeout") Integer fetchSequenceNumberTimeout, + @JsonProperty("fetchThreads") Integer fetchThreads, + @JsonProperty("segmentWriteOutMediumFactory") @Nullable SegmentWriteOutMediumFactory segmentWriteOutMediumFactory, + @JsonProperty("logParseExceptions") @Nullable Boolean logParseExceptions, + @JsonProperty("maxParseExceptions") @Nullable Integer maxParseExceptions, + @JsonProperty("maxSavedParseExceptions") @Nullable Integer maxSavedParseExceptions, + @JsonProperty("maxRecordsPerPoll") @Nullable Integer maxRecordsPerPoll, + @JsonProperty("intermediateHandoffPeriod") @Nullable Period intermediateHandoffPeriod + ) + { + super( + maxRowsInMemory, + maxBytesInMemory, + maxRowsPerSegment, + maxTotalRows, + intermediatePersistPeriod, + basePersistDirectory, + maxPendingPersists, + indexSpec, + true, + reportParseExceptions, + handoffConditionTimeout, + resetOffsetAutomatically, + skipSequenceNumberAvailabilityCheck, + segmentWriteOutMediumFactory, + intermediateHandoffPeriod, + logParseExceptions, + maxParseExceptions, + maxSavedParseExceptions + ); + this.recordBufferSize = recordBufferSize == null ? DEFAULT_RECORD_BUFFER_SIZE : recordBufferSize; + this.recordBufferOfferTimeout = recordBufferOfferTimeout == null + ? DEFAULT_RECORD_BUFFER_OFFER_TIMEOUT + : recordBufferOfferTimeout; + this.recordBufferFullWait = recordBufferFullWait == null ? DEFAULT_RECORD_BUFFER_FULL_WAIT : recordBufferFullWait; + this.fetchSequenceNumberTimeout = fetchSequenceNumberTimeout + == null ? DEFAULT_FETCH_SEQUENCE_NUMBER_TIMEOUT : fetchSequenceNumberTimeout; + this.fetchThreads = fetchThreads; // we handle this being null later + this.maxRecordsPerPoll = maxRecordsPerPoll == null ? DEFAULT_MAX_RECORDS_PER_POLL : maxRecordsPerPoll; + + Preconditions.checkArgument( + !(super.isResetOffsetAutomatically() && super.isSkipSequenceNumberAvailabilityCheck()), + "resetOffsetAutomatically cannot be used if skipSequenceNumberAvailabilityCheck=true" + ); + } + + @JsonProperty + public int getRecordBufferSize() + { + return recordBufferSize; + } + + @JsonProperty + public int getRecordBufferOfferTimeout() + { + return recordBufferOfferTimeout; + } + + @JsonProperty + public int getRecordBufferFullWait() + { + return recordBufferFullWait; + } + + @JsonProperty + public int getFetchSequenceNumberTimeout() + { + return fetchSequenceNumberTimeout; + } + + @JsonProperty + public Integer getFetchThreads() + { + return fetchThreads; + } + + @JsonProperty + public int getMaxRecordsPerPoll() + { + return maxRecordsPerPoll; + } + + @Override + public KinesisIndexTaskTuningConfig withBasePersistDirectory(File dir) + { + return new KinesisIndexTaskTuningConfig( + getMaxRowsInMemory(), + getMaxBytesInMemory(), + getMaxRowsPerSegment(), + getMaxTotalRows(), + getIntermediatePersistPeriod(), + dir, + getMaxPendingPersists(), + getIndexSpec(), + true, + isReportParseExceptions(), + getHandoffConditionTimeout(), + isResetOffsetAutomatically(), + isSkipSequenceNumberAvailabilityCheck(), + getRecordBufferSize(), + getRecordBufferOfferTimeout(), + getRecordBufferFullWait(), + getFetchSequenceNumberTimeout(), + getFetchThreads(), + getSegmentWriteOutMediumFactory(), + isLogParseExceptions(), + getMaxParseExceptions(), + getMaxSavedParseExceptions(), + getMaxRecordsPerPoll(), + getIntermediateHandoffPeriod() + ); + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + KinesisIndexTaskTuningConfig that = (KinesisIndexTaskTuningConfig) o; + return getMaxRowsInMemory() == that.getMaxRowsInMemory() && + getMaxBytesInMemory() == that.getMaxBytesInMemory() && + getMaxRowsPerSegment() == that.getMaxRowsPerSegment() && + getMaxPendingPersists() == that.getMaxPendingPersists() && + getBuildV9Directly() == that.getBuildV9Directly() && + isReportParseExceptions() == that.isReportParseExceptions() && + getHandoffConditionTimeout() == that.getHandoffConditionTimeout() && + isResetOffsetAutomatically() == that.isResetOffsetAutomatically() && + isSkipSequenceNumberAvailabilityCheck() == that.isSkipSequenceNumberAvailabilityCheck() && + getRecordBufferSize() == that.getRecordBufferSize() && + getRecordBufferOfferTimeout() == that.getRecordBufferOfferTimeout() && + getRecordBufferFullWait() == that.getRecordBufferFullWait() && + getFetchSequenceNumberTimeout() == that.getFetchSequenceNumberTimeout() && + isLogParseExceptions() == that.isLogParseExceptions() && + getMaxParseExceptions() == that.getMaxParseExceptions() && + getMaxSavedParseExceptions() == that.getMaxSavedParseExceptions() && + getMaxRecordsPerPoll() == that.getMaxRecordsPerPoll() && + Objects.equals(getIntermediatePersistPeriod(), that.getIntermediatePersistPeriod()) && + Objects.equals(getBasePersistDirectory(), that.getBasePersistDirectory()) && + Objects.equals(getIndexSpec(), that.getIndexSpec()) && + Objects.equals(getFetchThreads(), that.getFetchThreads()) && + Objects.equals(getSegmentWriteOutMediumFactory(), that.getSegmentWriteOutMediumFactory()) && + Objects.equals(getMaxTotalRows(), that.getMaxTotalRows()) && + Objects.equals(getIntermediateHandoffPeriod(), that.getIntermediateHandoffPeriod()); + } + + @Override + public int hashCode() + { + return Objects.hash( + getMaxRowsInMemory(), + getMaxBytesInMemory(), + getMaxRowsPerSegment(), + getMaxTotalRows(), + getIntermediatePersistPeriod(), + getBasePersistDirectory(), + getMaxPendingPersists(), + getIndexSpec(), + true, + isReportParseExceptions(), + getHandoffConditionTimeout(), + isResetOffsetAutomatically(), + isSkipSequenceNumberAvailabilityCheck(), + getRecordBufferSize(), + getRecordBufferOfferTimeout(), + getRecordBufferFullWait(), + getFetchSequenceNumberTimeout(), + getFetchThreads(), + getSegmentWriteOutMediumFactory(), + isLogParseExceptions(), + getMaxParseExceptions(), + getMaxSavedParseExceptions(), + getMaxRecordsPerPoll(), + getIntermediateHandoffPeriod() + ); + } + + @Override + public String toString() + { + return "KinesisIndexTaskTuningConfig{" + + "maxRowsInMemory=" + getMaxRowsInMemory() + + ", maxBytesInMemory=" + getMaxBytesInMemory() + + ", maxRowsPerSegment=" + getMaxRowsPerSegment() + + ", maxTotalRows=" + getMaxTotalRows() + + ", intermediatePersistPeriod=" + getIntermediatePersistPeriod() + + ", basePersistDirectory=" + getBasePersistDirectory() + + ", maxPendingPersists=" + getMaxPendingPersists() + + ", indexSpec=" + getIndexSpec() + + ", reportParseExceptions=" + isReportParseExceptions() + + ", handoffConditionTimeout=" + getHandoffConditionTimeout() + + ", resetOffsetAutomatically=" + isResetOffsetAutomatically() + + ", skipSequenceNumberAvailabilityCheck=" + isSkipSequenceNumberAvailabilityCheck() + + ", recordBufferSize=" + recordBufferSize + + ", recordBufferOfferTimeout=" + recordBufferOfferTimeout + + ", recordBufferFullWait=" + recordBufferFullWait + + ", fetchSequenceNumberTimeout=" + fetchSequenceNumberTimeout + + ", fetchThreads=" + fetchThreads + + ", segmentWriteOutMediumFactory=" + getSegmentWriteOutMediumFactory() + + ", logParseExceptions=" + isLogParseExceptions() + + ", maxParseExceptions=" + getMaxParseExceptions() + + ", maxSavedParseExceptions=" + getMaxSavedParseExceptions() + + ", maxRecordsPerPoll=" + maxRecordsPerPoll + + ", intermediateHandoffPeriod=" + getIntermediateHandoffPeriod() + + '}'; + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexingServiceModule.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexingServiceModule.java new file mode 100644 index 000000000000..cba5166f268d --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexingServiceModule.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.jsontype.NamedType; +import com.fasterxml.jackson.databind.module.SimpleModule; +import com.google.common.collect.ImmutableList; +import com.google.inject.Binder; +import com.google.inject.TypeLiteral; +import com.google.inject.name.Names; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.guice.JsonConfigProvider; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.indexing.kinesis.supervisor.KinesisSupervisorSpec; +import org.apache.druid.indexing.kinesis.supervisor.KinesisSupervisorTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; +import org.apache.druid.initialization.DruidModule; + +import java.util.List; + +public class KinesisIndexingServiceModule implements DruidModule +{ + @Override + public List getJacksonModules() + { + return ImmutableList.of( + new SimpleModule(getClass().getSimpleName()) + .registerSubtypes( + new NamedType(KinesisIndexTask.class, "index_kinesis"), + new NamedType(KinesisDataSourceMetadata.class, "kinesis"), + new NamedType(KinesisIndexTaskIOConfig.class, "kinesis"), + new NamedType(KinesisSupervisorTuningConfig.class, "kinesis"), + new NamedType(KinesisSupervisorSpec.class, "kinesis") + ) + ); + } + + @Override + public void configure(Binder binder) + { + binder.bind( + new TypeLiteral>() + { + } + ).to(KinesisIndexTaskClientFactory.class).in(LazySingleton.class); + + JsonConfigProvider.bind(binder, "druid.kinesis", AWSCredentialsConfig.class, Names.named("kinesis")); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplier.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplier.java new file mode 100644 index 000000000000..622c1f37bbc3 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplier.java @@ -0,0 +1,780 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.amazonaws.AmazonServiceException; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; +import com.amazonaws.client.builder.AwsClientBuilder; +import com.amazonaws.services.kinesis.AmazonKinesis; +import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder; +import com.amazonaws.services.kinesis.model.ExpiredIteratorException; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.InvalidArgumentException; +import com.amazonaws.services.kinesis.model.ProvisionedThroughputExceededException; +import com.amazonaws.services.kinesis.model.Record; +import com.amazonaws.services.kinesis.model.ResourceNotFoundException; +import com.amazonaws.services.kinesis.model.Shard; +import com.amazonaws.services.kinesis.model.ShardIteratorType; +import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; +import com.amazonaws.util.AwsHostNameUtils; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Queues; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.common.aws.AWSCredentialsUtils; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.emitter.EmittingLogger; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.io.IOException; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.reflect.Method; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +/** + * This class implements a local buffer for storing fetched Kinesis records. Fetching is done + * in background threads. + */ +public class KinesisRecordSupplier implements RecordSupplier +{ + private static final EmittingLogger log = new EmittingLogger(KinesisRecordSupplier.class); + private static final long PROVISIONED_THROUGHPUT_EXCEEDED_BACKOFF_MS = 3000; + private static final long EXCEPTION_RETRY_DELAY_MS = 10000; + + private static boolean isServiceExceptionRecoverable(AmazonServiceException ex) + { + final boolean isIOException = ex.getCause() instanceof IOException; + final boolean isTimeout = "RequestTimeout".equals(ex.getErrorCode()); + return isIOException || isTimeout; + } + + private class PartitionResource + { + private final StreamPartition streamPartition; + + // shardIterator points to the record that will be polled next by recordRunnable + // can be null when shard is closed due to the user shard splitting or changing the number + // of shards in the stream, in which case a 'EOS' marker is used by the KinesisRecordSupplier + // to indicate that this shard has no more records to read + @Nullable + private volatile String shardIterator; + private volatile boolean started; + private volatile boolean stopRequested; + + PartitionResource( + StreamPartition streamPartition + ) + { + this.streamPartition = streamPartition; + } + + void startBackgroundFetch() + { + if (started) { + return; + } + + log.info( + "Starting scheduled fetch runnable for stream[%s] partition[%s]", + streamPartition.getStream(), + streamPartition.getPartitionId() + ); + + stopRequested = false; + started = true; + + rescheduleRunnable(fetchDelayMillis); + } + + void stopBackgroundFetch() + { + log.info( + "Stopping scheduled fetch runnable for stream[%s] partition[%s]", + streamPartition.getStream(), + streamPartition.getPartitionId() + ); + stopRequested = true; + } + + + private Runnable getRecordRunnable() + { + return () -> { + + if (stopRequested) { + started = false; + stopRequested = false; + + log.info("Worker for partition[%s] has been stopped", streamPartition.getPartitionId()); + return; + } + + // used for retrying on InterruptedException + GetRecordsResult recordsResult = null; + OrderedPartitionableRecord currRecord; + + try { + + if (shardIterator == null) { + log.info("shardIterator[%s] has been closed and has no more records", streamPartition.getPartitionId()); + + // add an end-of-shard marker so caller knows this shard is closed + currRecord = new OrderedPartitionableRecord<>( + streamPartition.getStream(), + streamPartition.getPartitionId(), + KinesisSequenceNumber.END_OF_SHARD_MARKER, + null + ); + + recordsResult = null; + + if (!records.offer(currRecord, recordBufferOfferTimeout, TimeUnit.MILLISECONDS)) { + log.warn("OrderedPartitionableRecord buffer full, retrying in [%,dms]", recordBufferFullWait); + rescheduleRunnable(recordBufferFullWait); + } + + return; + } + + recordsResult = kinesis.getRecords(new GetRecordsRequest().withShardIterator( + shardIterator).withLimit(recordsPerFetch)); + + // list will come back empty if there are no records + for (Record kinesisRecord : recordsResult.getRecords()) { + + final List data; + + if (deaggregate) { + if (deaggregateHandle == null || getDataHandle == null) { + throw new ISE("deaggregateHandle or getDataHandle is null!"); + } + + data = new ArrayList<>(); + + final List userRecords = (List) deaggregateHandle.invokeExact( + Collections.singletonList(kinesisRecord) + ); + + for (Object userRecord : userRecords) { + data.add(toByteArray((ByteBuffer) getDataHandle.invoke(userRecord))); + } + } else { + data = Collections.singletonList(toByteArray(kinesisRecord.getData())); + } + + currRecord = new OrderedPartitionableRecord<>( + streamPartition.getStream(), + streamPartition.getPartitionId(), + kinesisRecord.getSequenceNumber(), + data + ); + + + log.trace( + "Stream[%s] / partition[%s] / sequenceNum[%s] / bufferRemainingCapacity[%d]: %s", + currRecord.getStream(), + currRecord.getPartitionId(), + currRecord.getSequenceNumber(), + records.remainingCapacity(), + currRecord.getData().stream().map(StringUtils::fromUtf8).collect(Collectors.toList()) + ); + + // If the buffer was full and we weren't able to add the message, grab a new stream iterator starting + // from this message and back off for a bit to let the buffer drain before retrying. + if (!records.offer(currRecord, recordBufferOfferTimeout, TimeUnit.MILLISECONDS)) { + log.warn( + "OrderedPartitionableRecord buffer full, storing iterator and retrying in [%,dms]", + recordBufferFullWait + ); + + shardIterator = kinesis.getShardIterator( + currRecord.getStream(), + currRecord.getPartitionId(), + ShardIteratorType.AT_SEQUENCE_NUMBER.toString(), + currRecord.getSequenceNumber() + ).getShardIterator(); + + rescheduleRunnable(recordBufferFullWait); + return; + } + } + + shardIterator = recordsResult.getNextShardIterator(); // will be null if the shard has been closed + + rescheduleRunnable(fetchDelayMillis); + } + catch (ProvisionedThroughputExceededException e) { + log.warn( + e, + "encounted ProvisionedThroughputExceededException while fetching records, this means " + + "that the request rate for the stream is too high, or the requested data is too large for " + + "the available throughput. Reduce the frequency or size of your requests." + ); + long retryMs = Math.max(PROVISIONED_THROUGHPUT_EXCEEDED_BACKOFF_MS, fetchDelayMillis); + rescheduleRunnable(retryMs); + } + catch (InterruptedException e) { + // may happen if interrupted while BlockingQueue.offer() is waiting + log.warn( + e, + "Interrupted while waiting to add record to buffer, retrying in [%,dms]", + EXCEPTION_RETRY_DELAY_MS + ); + rescheduleRunnable(EXCEPTION_RETRY_DELAY_MS); + } + catch (ExpiredIteratorException e) { + log.warn( + e, + "ShardIterator expired while trying to fetch records, retrying in [%,dms]", + fetchDelayMillis + ); + if (recordsResult != null) { + shardIterator = recordsResult.getNextShardIterator(); // will be null if the shard has been closed + rescheduleRunnable(fetchDelayMillis); + } else { + throw new ISE("can't reschedule fetch records runnable, recordsResult is null??"); + } + } + catch (ResourceNotFoundException | InvalidArgumentException e) { + // aws errors + log.error(e, "encounted AWS error while attempting to fetch records, will not retry"); + throw e; + } + catch (AmazonServiceException e) { + if (isServiceExceptionRecoverable(e)) { + log.warn(e, "encounted unknown recoverable AWS exception, retrying in [%,dms]", EXCEPTION_RETRY_DELAY_MS); + rescheduleRunnable(EXCEPTION_RETRY_DELAY_MS); + } else { + log.warn(e, "encounted unknown unrecoverable AWS exception, will not retry"); + throw new RuntimeException(e); + } + } + catch (Throwable e) { + // non transient errors + log.error(e, "unknown getRecordRunnable exception, will not retry"); + throw new RuntimeException(e); + } + + }; + } + + private void rescheduleRunnable(long delayMillis) + { + if (started && !stopRequested) { + try { + scheduledExec.schedule(getRecordRunnable(), delayMillis, TimeUnit.MILLISECONDS); + } + catch (RejectedExecutionException e) { + log.warn( + e, + "Caught RejectedExecutionException, KinesisRecordSupplier for partition[%s] has likely temporarily shutdown the ExecutorService." + + "This is expected behavior after calling seek(), seekToEarliest() and seekToLatest()", + streamPartition.getPartitionId() + ); + + } + } else { + log.info("Worker for partition[%s] has been stopped", streamPartition.getPartitionId()); + } + } + } + + // used for deaggregate + private final MethodHandle deaggregateHandle; + private final MethodHandle getDataHandle; + + private final AmazonKinesis kinesis; + + private final int recordsPerFetch; + private final int fetchDelayMillis; + private final boolean deaggregate; + private final int recordBufferOfferTimeout; + private final int recordBufferFullWait; + private final int fetchSequenceNumberTimeout; + private final int maxRecordsPerPoll; + private final int fetchThreads; + private final int recordBufferSize; + + private ScheduledExecutorService scheduledExec; + + private final Map, PartitionResource> partitionResources = new ConcurrentHashMap<>(); + private BlockingQueue> records; + + private volatile boolean checkPartitionsStarted = false; + private volatile boolean closed = false; + + public KinesisRecordSupplier( + AmazonKinesis amazonKinesis, + int recordsPerFetch, + int fetchDelayMillis, + int fetchThreads, + boolean deaggregate, + int recordBufferSize, + int recordBufferOfferTimeout, + int recordBufferFullWait, + int fetchSequenceNumberTimeout, + int maxRecordsPerPoll + ) + { + Preconditions.checkNotNull(amazonKinesis); + this.kinesis = amazonKinesis; + this.recordsPerFetch = recordsPerFetch; + this.fetchDelayMillis = fetchDelayMillis; + this.deaggregate = deaggregate; + this.recordBufferOfferTimeout = recordBufferOfferTimeout; + this.recordBufferFullWait = recordBufferFullWait; + this.fetchSequenceNumberTimeout = fetchSequenceNumberTimeout; + this.maxRecordsPerPoll = maxRecordsPerPoll; + this.fetchThreads = fetchThreads; + this.recordBufferSize = recordBufferSize; + + // the deaggregate function is implemented by the amazon-kinesis-client, whose license is not compatible with Apache. + // The work around here is to use reflection to find the deaggregate function in the classpath. See details on the + // docs page for more information on how to use deaggregation + if (deaggregate) { + try { + Class kclUserRecordclass = Class.forName("com.amazonaws.services.kinesis.clientlibrary.types.UserRecord"); + MethodHandles.Lookup lookup = MethodHandles.publicLookup(); + + Method deaggregateMethod = kclUserRecordclass.getMethod("deaggregate", List.class); + Method getDataMethod = kclUserRecordclass.getMethod("getData"); + + deaggregateHandle = lookup.unreflect(deaggregateMethod); + getDataHandle = lookup.unreflect(getDataMethod); + } + catch (ClassNotFoundException e) { + throw new ISE(e, "cannot find class[com.amazonaws.services.kinesis.clientlibrary.types.UserRecord], " + + "note that when using deaggregate=true, you must provide the Kinesis Client Library jar in the classpath"); + } + catch (Exception e) { + throw new RuntimeException(e); + } + } else { + deaggregateHandle = null; + getDataHandle = null; + } + + log.info( + "Creating fetch thread pool of size [%d] (Runtime.availableProcessors=%d)", + fetchThreads, + Runtime.getRuntime().availableProcessors() + ); + + scheduledExec = Executors.newScheduledThreadPool( + fetchThreads, + Execs.makeThreadFactory("KinesisRecordSupplier-Worker-%d") + ); + + records = new LinkedBlockingQueue<>(recordBufferSize); + } + + public static AmazonKinesis getAmazonKinesisClient( + String endpoint, + AWSCredentialsConfig awsCredentialsConfig, + String awsAssumedRoleArn, + String awsExternalId + ) + { + AWSCredentialsProvider awsCredentialsProvider = AWSCredentialsUtils.defaultAWSCredentialsProviderChain( + awsCredentialsConfig + ); + + if (awsAssumedRoleArn != null) { + log.info("Assuming role [%s] with externalId [%s]", awsAssumedRoleArn, awsExternalId); + + STSAssumeRoleSessionCredentialsProvider.Builder builder = new STSAssumeRoleSessionCredentialsProvider + .Builder(awsAssumedRoleArn, StringUtils.format("druid-kinesis-%s", UUID.randomUUID().toString())) + .withStsClient(AWSSecurityTokenServiceClientBuilder.standard() + .withCredentials(awsCredentialsProvider) + .build()); + + if (awsExternalId != null) { + builder.withExternalId(awsExternalId); + } + + awsCredentialsProvider = builder.build(); + } + + return AmazonKinesisClientBuilder.standard() + .withCredentials(awsCredentialsProvider) + .withClientConfiguration(new ClientConfiguration()) + .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( + endpoint, + AwsHostNameUtils.parseRegion( + endpoint, + null + ) + )).build(); + } + + + @VisibleForTesting + public void start() + { + checkIfClosed(); + if (checkPartitionsStarted) { + partitionResources.values().forEach(PartitionResource::startBackgroundFetch); + checkPartitionsStarted = false; + } + } + + @Override + public void assign(Set> collection) + { + checkIfClosed(); + + collection.forEach( + streamPartition -> partitionResources.putIfAbsent( + streamPartition, + new PartitionResource(streamPartition) + ) + ); + + for (Iterator, PartitionResource>> i = partitionResources.entrySet() + .iterator(); i.hasNext(); ) { + Map.Entry, PartitionResource> entry = i.next(); + if (!collection.contains(entry.getKey())) { + i.remove(); + entry.getValue().stopBackgroundFetch(); + } + } + + } + + @Override + public void seek(StreamPartition partition, String sequenceNumber) throws InterruptedException + { + checkIfClosed(); + filterBufferAndResetFetchRunnable(ImmutableSet.of(partition)); + seekInternal(partition, sequenceNumber, ShardIteratorType.AT_SEQUENCE_NUMBER); + } + + @Override + public void seekToEarliest(Set> partitions) throws InterruptedException + { + checkIfClosed(); + filterBufferAndResetFetchRunnable(partitions); + partitions.forEach(partition -> seekInternal(partition, null, ShardIteratorType.TRIM_HORIZON)); + } + + @Override + public void seekToLatest(Set> partitions) throws InterruptedException + { + checkIfClosed(); + filterBufferAndResetFetchRunnable(partitions); + partitions.forEach(partition -> seekInternal(partition, null, ShardIteratorType.LATEST)); + } + + @Override + public Collection> getAssignment() + { + return partitionResources.keySet(); + } + + @Nonnull + @Override + public List> poll(long timeout) + { + checkIfClosed(); + if (checkPartitionsStarted) { + partitionResources.values().forEach(PartitionResource::startBackgroundFetch); + checkPartitionsStarted = false; + } + + try { + int expectedSize = Math.min(Math.max(records.size(), 1), maxRecordsPerPoll); + + List> polledRecords = new ArrayList<>(expectedSize); + + Queues.drain( + records, + polledRecords, + expectedSize, + timeout, + TimeUnit.MILLISECONDS + ); + + polledRecords = polledRecords.stream() + .filter(x -> partitionResources.containsKey(x.getStreamPartition())) + .collect(Collectors.toList()); + + return polledRecords; + } + catch (InterruptedException e) { + log.warn(e, "Interrupted while polling"); + return Collections.emptyList(); + } + + } + + @Nullable + @Override + public String getLatestSequenceNumber(StreamPartition partition) + { + checkIfClosed(); + return getSequenceNumberInternal(partition, ShardIteratorType.LATEST); + } + + @Nullable + @Override + public String getEarliestSequenceNumber(StreamPartition partition) + { + checkIfClosed(); + return getSequenceNumberInternal(partition, ShardIteratorType.TRIM_HORIZON); + } + + @Nullable + @Override + public String getPosition(StreamPartition partition) + { + throw new UnsupportedOperationException("getPosition() is not supported in Kinesis"); + } + + @Override + public Set getPartitionIds(String stream) + { + checkIfClosed(); + return kinesis.describeStream(stream) + .getStreamDescription() + .getShards() + .stream() + .map(Shard::getShardId).collect(Collectors.toSet()); + } + + @Override + public void close() + { + if (this.closed) { + return; + } + + assign(ImmutableSet.of()); + + scheduledExec.shutdown(); + + try { + if (!scheduledExec.awaitTermination(EXCEPTION_RETRY_DELAY_MS, TimeUnit.MILLISECONDS)) { + scheduledExec.shutdownNow(); + } + } + catch (InterruptedException e) { + log.warn(e, "InterruptedException while shutting down"); + throw new RuntimeException(e); + } + + this.closed = true; + } + + private void seekInternal(StreamPartition partition, String sequenceNumber, ShardIteratorType iteratorEnum) + { + PartitionResource resource = partitionResources.get(partition); + if (resource == null) { + throw new ISE("Partition [%s] has not been assigned", partition); + } + + log.debug( + "Seeking partition [%s] to [%s]", + partition.getPartitionId(), + sequenceNumber != null ? sequenceNumber : iteratorEnum.toString() + ); + + resource.shardIterator = kinesis.getShardIterator( + partition.getStream(), + partition.getPartitionId(), + iteratorEnum.toString(), + sequenceNumber + ).getShardIterator(); + + checkPartitionsStarted = true; + } + + private void filterBufferAndResetFetchRunnable(Set> partitions) throws InterruptedException + { + scheduledExec.shutdown(); + + try { + if (!scheduledExec.awaitTermination(EXCEPTION_RETRY_DELAY_MS, TimeUnit.MILLISECONDS)) { + scheduledExec.shutdownNow(); + } + } + catch (InterruptedException e) { + log.warn(e, "InterruptedException while shutting down"); + throw e; + } + + scheduledExec = Executors.newScheduledThreadPool( + fetchThreads, + Execs.makeThreadFactory("KinesisRecordSupplier-Worker-%d") + ); + + // filter records in buffer and only retain ones whose partition was not seeked + BlockingQueue> newQ = new LinkedBlockingQueue<>(recordBufferSize); + records + .stream() + .filter(x -> !partitions.contains(x.getStreamPartition())) + .forEachOrdered(newQ::offer); + + records = newQ; + + // restart fetching threads + partitionResources.values().forEach(x -> x.started = false); + checkPartitionsStarted = true; + } + + @Nullable + private String getSequenceNumberInternal(StreamPartition partition, ShardIteratorType iteratorEnum) + { + + String shardIterator = null; + try { + shardIterator = kinesis.getShardIterator( + partition.getStream(), + partition.getPartitionId(), + iteratorEnum.toString() + ).getShardIterator(); + } + catch (ResourceNotFoundException e) { + log.warn(e, "Caught ResourceNotFoundException while getting shardIterator"); + } + + return getSequenceNumberInternal(partition, shardIterator); + } + + @Nullable + private String getSequenceNumberInternal(StreamPartition partition, String shardIterator) + { + long timeoutMillis = System.currentTimeMillis() + fetchSequenceNumberTimeout; + + while (shardIterator != null && System.currentTimeMillis() < timeoutMillis) { + + if (closed) { + log.info("KinesisRecordSupplier closed while fetching sequenceNumber"); + return null; + } + + GetRecordsResult recordsResult; + try { + // we call getRecords with limit 1000 to make sure that we can find the first (earliest) record in the shard. + // In the case where the shard is constantly removing records that are past their retention period, it is possible + // that we never find the first record in the shard if we use a limit of 1. + recordsResult = kinesis.getRecords(new GetRecordsRequest().withShardIterator(shardIterator).withLimit(1000)); + } + catch (ProvisionedThroughputExceededException e) { + log.warn( + e, + "encountered ProvisionedThroughputExceededException while fetching records, this means " + + "that the request rate for the stream is too high, or the requested data is too large for " + + "the available throughput. Reduce the frequency or size of your requests. Consider increasing " + + "the number of shards to increase throughput." + ); + try { + Thread.sleep(PROVISIONED_THROUGHPUT_EXCEEDED_BACKOFF_MS); + continue; + } + catch (InterruptedException e1) { + log.warn(e1, "Thread interrupted!"); + Thread.currentThread().interrupt(); + break; + } + } + + List records = recordsResult.getRecords(); + + if (!records.isEmpty()) { + return records.get(0).getSequenceNumber(); + } + + shardIterator = recordsResult.getNextShardIterator(); + } + + if (shardIterator == null) { + log.info("Partition[%s] returned a null shard iterator, is the shard closed?", partition.getPartitionId()); + return KinesisSequenceNumber.END_OF_SHARD_MARKER; + } + + + // if we reach here, it usually means either the shard has no more records, or records have not been + // added to this shard + log.warn( + "timed out while trying to fetch position for shard[%s], likely no more records in shard", + partition.getPartitionId() + ); + return null; + + } + + private void checkIfClosed() + { + if (closed) { + throw new ISE("Invalid operation - KinesisRecordSupplier has already been closed"); + } + } + + /** + * Returns an array with the content between the position and limit of "buffer". This may be the buffer's backing + * array itself. Does not modify position or limit of the buffer. + */ + private static byte[] toByteArray(final ByteBuffer buffer) + { + if (buffer.hasArray() + && buffer.arrayOffset() == 0 + && buffer.position() == 0 + && buffer.array().length == buffer.limit()) { + return buffer.array(); + } else { + final byte[] retVal = new byte[buffer.remaining()]; + buffer.duplicate().get(retVal); + return retVal; + } + } + + @VisibleForTesting + public int bufferSize() + { + return records.size(); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRegion.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRegion.java new file mode 100644 index 000000000000..a5e1195616f9 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisRegion.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; +import org.apache.druid.java.util.common.IAE; +import org.apache.druid.java.util.common.StringUtils; + +import java.util.EnumSet; +import java.util.List; +import java.util.stream.Collectors; + +public enum KinesisRegion +{ + US_EAST_2, + US_EAST_1, + US_WEST_1, + US_WEST_2, + AP_NORTHEAST_1, + AP_NORTHEAST_2, + AP_NORTHEAST_3, + AP_SOUTH_1, + AP_SOUTHEAST_1, + AP_SOUTHEAST_2, + CA_CENTRAL_1, + CN_NORTH_1, + CN_NORTHWEST_1, + EU_CENTRAL_1, + EU_WEST_1, + EU_WEST_2, + EU_WEST_3, + SA_EAST_1, + US_GOV_WEST_1; + + @JsonCreator + public static KinesisRegion fromString(String value) + { + return EnumSet.allOf(KinesisRegion.class) + .stream() + .filter(x -> x.serialize().equals(value)) + .findFirst() + .orElseThrow(() -> new IAE("Invalid region %s, region must be one of: %s", value, getNames())); + } + + private static List getNames() + { + return EnumSet.allOf(KinesisRegion.class).stream().map(KinesisRegion::serialize).collect(Collectors.toList()); + } + + public String getEndpoint() + { + return StringUtils.format("kinesis.%s.amazonaws.com%s", serialize(), serialize().startsWith("cn-") ? ".cn" : ""); + } + + @JsonValue + public String serialize() + { + return StringUtils.toLowerCase(name()).replace('_', '-'); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisSequenceNumber.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisSequenceNumber.java new file mode 100644 index 000000000000..9a96e4ed5985 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisSequenceNumber.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + + +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; + +import javax.validation.constraints.NotNull; +import java.math.BigInteger; + +public class KinesisSequenceNumber extends OrderedSequenceNumber +{ + + /** + * In Kinesis, when a shard is closed due to shard splitting, a null ShardIterator is returned. + * The EOS marker is placed at the end of the Kinesis Record Supplier buffer, such that when + * an indexing task pulls the record 'EOS', it knows the shard has been closed and should stop + * reading and start publishing + */ + public static final String END_OF_SHARD_MARKER = "EOS"; + // this flag is used to indicate either END_OF_SHARD_MARKER + // or NO_END_SEQUENCE_NUMBER so that they can be properly compared + // with other sequence numbers + private final boolean isMaxSequenceNumber; + private final BigInteger intSequence; + + private KinesisSequenceNumber(@NotNull String sequenceNumber, boolean isExclusive) + { + super(sequenceNumber, isExclusive); + if (END_OF_SHARD_MARKER.equals(sequenceNumber) + || SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER.equals(sequenceNumber)) { + isMaxSequenceNumber = true; + this.intSequence = null; + } else { + isMaxSequenceNumber = false; + this.intSequence = new BigInteger(sequenceNumber); + } + } + + public static KinesisSequenceNumber of(String sequenceNumber) + { + return new KinesisSequenceNumber(sequenceNumber, false); + } + + public static KinesisSequenceNumber of(String sequenceNumber, boolean isExclusive) + { + return new KinesisSequenceNumber(sequenceNumber, isExclusive); + } + + @Override + public int compareTo(@NotNull OrderedSequenceNumber o) + { + KinesisSequenceNumber num = (KinesisSequenceNumber) o; + if (isMaxSequenceNumber && num.isMaxSequenceNumber) { + return 0; + } else if (isMaxSequenceNumber) { + return 1; + } else if (num.isMaxSequenceNumber) { + return -1; + } else { + return this.intSequence.compareTo(new BigInteger(o.get())); + } + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisor.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisor.java new file mode 100644 index 000000000000..f53bdf2ad129 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisor.java @@ -0,0 +1,311 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.common.task.TaskResource; +import org.apache.druid.indexing.kinesis.KinesisDataSourceMetadata; +import org.apache.druid.indexing.kinesis.KinesisIndexTask; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskClientFactory; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskIOConfig; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskTuningConfig; +import org.apache.druid.indexing.kinesis.KinesisRecordSupplier; +import org.apache.druid.indexing.kinesis.KinesisSequenceNumber; +import org.apache.druid.indexing.overlord.DataSourceMetadata; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorReportPayload; +import org.apache.druid.indexing.seekablestream.utils.RandomIdUtils; +import org.apache.druid.java.util.common.StringUtils; +import org.joda.time.DateTime; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.ScheduledExecutorService; + +/** + * Supervisor responsible for managing the KinesisIndexTask for a single dataSource. At a high level, the class accepts a + * {@link KinesisSupervisorSpec} which includes the Kinesis stream and configuration as well as an ingestion spec which will + * be used to generate the indexing tasks. The run loop periodically refreshes its view of the Kinesis stream's partitions + * and the list of running indexing tasks and ensures that all partitions are being read from and that there are enough + * tasks to satisfy the desired number of replicas. As tasks complete, new tasks are queued to process the next range of + * Kinesis sequences. + *

+ * the Kinesis supervisor does not yet support lag calculations + */ +public class KinesisSupervisor extends SeekableStreamSupervisor +{ + private static final String NOT_SET = "-1"; + private final KinesisSupervisorSpec spec; + private final AWSCredentialsConfig awsCredentialsConfig; + + public KinesisSupervisor( + final TaskStorage taskStorage, + final TaskMaster taskMaster, + final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + final KinesisIndexTaskClientFactory taskClientFactory, + final ObjectMapper mapper, + final KinesisSupervisorSpec spec, + final RowIngestionMetersFactory rowIngestionMetersFactory, + final AWSCredentialsConfig awsCredentialsConfig + ) + { + super( + StringUtils.format("KinesisSupervisor-%s", spec.getDataSchema().getDataSource()), + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + mapper, + spec, + rowIngestionMetersFactory, + true + ); + + this.spec = spec; + this.awsCredentialsConfig = awsCredentialsConfig; + } + + @Override + protected SeekableStreamIndexTaskIOConfig createTaskIoConfig( + int groupId, + Map startPartitions, + Map endPartitions, + String baseSequenceName, + DateTime minimumMessageTime, + DateTime maximumMessageTime, + Set exclusiveStartSequenceNumberPartitions, + SeekableStreamSupervisorIOConfig ioConfigg + ) + { + KinesisSupervisorIOConfig ioConfig = (KinesisSupervisorIOConfig) ioConfigg; + return new KinesisIndexTaskIOConfig( + groupId, + baseSequenceName, + new SeekableStreamPartitions<>(ioConfig.getStream(), startPartitions), + new SeekableStreamPartitions<>(ioConfig.getStream(), endPartitions), + true, + minimumMessageTime, + maximumMessageTime, + ioConfig.getEndpoint(), + ioConfig.getRecordsPerFetch(), + ioConfig.getFetchDelayMillis(), + exclusiveStartSequenceNumberPartitions, + ioConfig.getAwsAssumedRoleArn(), + ioConfig.getAwsExternalId(), + ioConfig.isDeaggregate() + ); + } + + @Override + protected List> createIndexTasks( + int replicas, + String baseSequenceName, + ObjectMapper sortingMapper, + TreeMap> sequenceOffsets, + SeekableStreamIndexTaskIOConfig taskIoConfig, + SeekableStreamIndexTaskTuningConfig taskTuningConfig, + RowIngestionMetersFactory rowIngestionMetersFactory + ) throws JsonProcessingException + { + final String checkpoints = sortingMapper.writerFor(new TypeReference>>() + { + }).writeValueAsString(sequenceOffsets); + final Map context = spec.getContext() == null + ? ImmutableMap.of( + "checkpoints", + checkpoints, + IS_INCREMENTAL_HANDOFF_SUPPORTED, + true + ) : ImmutableMap.builder() + .put("checkpoints", checkpoints) + .put(IS_INCREMENTAL_HANDOFF_SUPPORTED, true) + .putAll(spec.getContext()) + .build(); + List> taskList = new ArrayList<>(); + for (int i = 0; i < replicas; i++) { + String taskId = Joiner.on("_").join(baseSequenceName, RandomIdUtils.getRandomId()); + taskList.add(new KinesisIndexTask( + taskId, + new TaskResource(baseSequenceName, 1), + spec.getDataSchema(), + (KinesisIndexTaskTuningConfig) taskTuningConfig, + (KinesisIndexTaskIOConfig) taskIoConfig, + context, + null, + null, + rowIngestionMetersFactory, + awsCredentialsConfig + )); + } + return taskList; + } + + + @Override + protected RecordSupplier setupRecordSupplier() + throws RuntimeException + { + KinesisSupervisorIOConfig ioConfig = spec.getIoConfig(); + KinesisIndexTaskTuningConfig taskTuningConfig = spec.getTuningConfig(); + + return new KinesisRecordSupplier( + KinesisRecordSupplier.getAmazonKinesisClient( + ioConfig.getEndpoint(), + awsCredentialsConfig, + ioConfig.getAwsAssumedRoleArn(), + ioConfig.getAwsExternalId() + ), + ioConfig.getRecordsPerFetch(), + ioConfig.getFetchDelayMillis(), + 1, + ioConfig.isDeaggregate(), + taskTuningConfig.getRecordBufferSize(), + taskTuningConfig.getRecordBufferOfferTimeout(), + taskTuningConfig.getRecordBufferFullWait(), + taskTuningConfig.getFetchSequenceNumberTimeout(), + taskTuningConfig.getMaxRecordsPerPoll() + ); + + } + + + @Override + protected void scheduleReporting(ScheduledExecutorService reportingExec) + { + // not yet implemented, see issue #6739 + } + + @Override + protected int getTaskGroupIdForPartition(String partitionId) + { + if (!partitionIds.contains(partitionId)) { + partitionIds.add(partitionId); + } + + return partitionIds.indexOf(partitionId) % spec.getIoConfig().getTaskCount(); + } + + @Override + protected boolean checkSourceMetadataMatch(DataSourceMetadata metadata) + { + return metadata instanceof KinesisDataSourceMetadata; + } + + @Override + protected boolean doesTaskTypeMatchSupervisor(Task task) + { + return task instanceof KinesisIndexTask; + } + + @Override + protected SeekableStreamSupervisorReportPayload createReportPayload( + int numPartitions, + boolean includeOffsets + ) + { + KinesisSupervisorIOConfig ioConfig = spec.getIoConfig(); + return new KinesisSupervisorReportPayload( + spec.getDataSchema().getDataSource(), + ioConfig.getStream(), + numPartitions, + ioConfig.getReplicas(), + ioConfig.getTaskDuration().getMillis() / 1000, + spec.isSuspended() + ); + } + + // not yet supported, will be implemented in the future + @Override + protected Map getLagPerPartition(Map currentOffsets) + { + return ImmutableMap.of(); + } + + @Override + protected SeekableStreamDataSourceMetadata createDataSourceMetaData( + String stream, + Map map + ) + { + return new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>(stream, map) + ); + } + + @Override + protected OrderedSequenceNumber makeSequenceNumber(String seq, boolean isExclusive) + { + return KinesisSequenceNumber.of(seq, isExclusive); + } + + @Override + protected void updateLatestSequenceFromStream( + RecordSupplier recordSupplier, Set> streamPartitions + ) + { + // do nothing + } + + @Override + protected String baseTaskName() + { + return "index_kinesis"; + } + + @Override + protected String getNotSetMarker() + { + return NOT_SET; + } + + @Override + protected String getEndOfPartitionMarker() + { + return SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER; + } + + @Override + protected boolean isEndOfShard(String seqNum) + { + return KinesisSequenceNumber.END_OF_SHARD_MARKER.equals(seqNum); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfig.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfig.java new file mode 100644 index 000000000000..2fb43d315175 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfig.java @@ -0,0 +1,156 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskIOConfig; +import org.apache.druid.indexing.kinesis.KinesisRegion; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig; +import org.joda.time.Period; + +public class KinesisSupervisorIOConfig extends SeekableStreamSupervisorIOConfig +{ + private final String endpoint; + + // In determining a suitable value for recordsPerFetch: + // - Each data record can be up to 1 MB in size + // - Each shard can read up to 2 MB per second + // - The maximum size of data that GetRecords can return is 10 MB. If a call returns this amount of data, + // subsequent calls made within the next 5 seconds throw ProvisionedThroughputExceededException. + // + // If there is insufficient provisioned throughput on the shard, subsequent calls made within the next 1 second + // throw ProvisionedThroughputExceededException. Note that GetRecords won't return any data when it throws an + // exception. For this reason, we recommend that you wait one second between calls to GetRecords; however, it's + // possible that the application will get exceptions for longer than 1 second. + private final Integer recordsPerFetch; + private final Integer fetchDelayMillis; + + private final String awsAssumedRoleArn; + private final String awsExternalId; + private final boolean deaggregate; + + @JsonCreator + public KinesisSupervisorIOConfig( + @JsonProperty("stream") String stream, + @JsonProperty("endpoint") String endpoint, + @JsonProperty("region") KinesisRegion region, + @JsonProperty("replicas") Integer replicas, + @JsonProperty("taskCount") Integer taskCount, + @JsonProperty("taskDuration") Period taskDuration, + @JsonProperty("startDelay") Period startDelay, + @JsonProperty("period") Period period, + @JsonProperty("useEarliestSequenceNumber") Boolean useEarliestSequenceNumber, + @JsonProperty("completionTimeout") Period completionTimeout, + @JsonProperty("lateMessageRejectionPeriod") Period lateMessageRejectionPeriod, + @JsonProperty("earlyMessageRejectionPeriod") Period earlyMessageRejectionPeriod, + @JsonProperty("recordsPerFetch") Integer recordsPerFetch, + @JsonProperty("fetchDelayMillis") Integer fetchDelayMillis, + @JsonProperty("awsAssumedRoleArn") String awsAssumedRoleArn, + @JsonProperty("awsExternalId") String awsExternalId, + @JsonProperty("deaggregate") boolean deaggregate + ) + { + super( + Preconditions.checkNotNull(stream, "stream"), + replicas, + taskCount, + taskDuration, + startDelay, + period, + useEarliestSequenceNumber, + completionTimeout, + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod + ); + this.endpoint = endpoint != null + ? endpoint + : (region != null ? region.getEndpoint() : KinesisRegion.US_EAST_1.getEndpoint()); + this.recordsPerFetch = recordsPerFetch != null + ? recordsPerFetch + : KinesisIndexTaskIOConfig.DEFAULT_RECORDS_PER_FETCH; + this.fetchDelayMillis = fetchDelayMillis != null + ? fetchDelayMillis + : KinesisIndexTaskIOConfig.DEFAULT_FETCH_DELAY_MILLIS; + this.awsAssumedRoleArn = awsAssumedRoleArn; + this.awsExternalId = awsExternalId; + this.deaggregate = deaggregate; + } + + @JsonProperty + public String getEndpoint() + { + return endpoint; + } + + @JsonProperty + public Integer getRecordsPerFetch() + { + return recordsPerFetch; + } + + @JsonProperty + public Integer getFetchDelayMillis() + { + return fetchDelayMillis; + } + + @JsonProperty + public String getAwsAssumedRoleArn() + { + return awsAssumedRoleArn; + } + + @JsonProperty + public String getAwsExternalId() + { + return awsExternalId; + } + + @JsonProperty + public boolean isDeaggregate() + { + return deaggregate; + } + + @Override + public String toString() + { + return "KinesisSupervisorIOConfig{" + + "stream='" + getStream() + '\'' + + ", endpoint='" + endpoint + '\'' + + ", replicas=" + getReplicas() + + ", taskCount=" + getTaskCount() + + ", taskDuration=" + getTaskDuration() + + ", startDelay=" + getStartDelay() + + ", period=" + getPeriod() + + ", useEarliestSequenceNumber=" + isUseEarliestSequenceNumber() + + ", completionTimeout=" + getCompletionTimeout() + + ", lateMessageRejectionPeriod=" + getLateMessageRejectionPeriod() + + ", earlyMessageRejectionPeriod=" + getEarlyMessageRejectionPeriod() + + ", recordsPerFetch=" + recordsPerFetch + + ", fetchDelayMillis=" + fetchDelayMillis + + ", awsAssumedRoleArn='" + awsAssumedRoleArn + '\'' + + ", awsExternalId='" + awsExternalId + '\'' + + ", deaggregate=" + deaggregate + + '}'; + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorReportPayload.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorReportPayload.java new file mode 100644 index 000000000000..fb08337b4642 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorReportPayload.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorReportPayload; + +import java.util.Collections; + +public class KinesisSupervisorReportPayload extends SeekableStreamSupervisorReportPayload +{ + public KinesisSupervisorReportPayload( + String dataSource, + String stream, + Integer partitions, + Integer replicas, + Long durationSeconds, + boolean suspended + ) + { + super( + dataSource, + stream, + partitions, + replicas, + durationSeconds, + Collections.emptyMap(), + Collections.emptyMap(), + null, + null, + suspended + ); + } + + @Override + public String toString() + { + return "KinesisSupervisorReportPayload{" + + "dataSource='" + getDataSource() + '\'' + + ", stream='" + getStream() + '\'' + + ", partitions=" + getPartitions() + + ", replicas=" + getReplicas() + + ", durationSeconds=" + getDurationSeconds() + + ", active=" + getActiveTasks() + + ", publishing=" + getPublishingTasks() + + ", suspended=" + getSuspended() + + '}'; + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorSpec.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorSpec.java new file mode 100644 index 000000000000..479851f8eb8c --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorSpec.java @@ -0,0 +1,176 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.annotation.JacksonInject; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.name.Named; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskClientFactory; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.supervisor.Supervisor; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorSpec; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.server.metrics.DruidMonitorSchedulerConfig; + +import java.util.Map; + +public class KinesisSupervisorSpec extends SeekableStreamSupervisorSpec +{ + private final AWSCredentialsConfig awsCredentialsConfig; + + @JsonCreator + public KinesisSupervisorSpec( + @JsonProperty("dataSchema") DataSchema dataSchema, + @JsonProperty("tuningConfig") KinesisSupervisorTuningConfig tuningConfig, + @JsonProperty("ioConfig") KinesisSupervisorIOConfig ioConfig, + @JsonProperty("context") Map context, + @JsonProperty("suspended") Boolean suspended, + @JacksonInject TaskStorage taskStorage, + @JacksonInject TaskMaster taskMaster, + @JacksonInject IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + @JacksonInject KinesisIndexTaskClientFactory kinesisIndexTaskClientFactory, + @JacksonInject @Json ObjectMapper mapper, + @JacksonInject ServiceEmitter emitter, + @JacksonInject DruidMonitorSchedulerConfig monitorSchedulerConfig, + @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory, + @JacksonInject @Named("kinesis") AWSCredentialsConfig awsCredentialsConfig + ) + { + super( + dataSchema, + tuningConfig != null + ? tuningConfig + : new KinesisSupervisorTuningConfig( + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ), + ioConfig, + context, + suspended, + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + kinesisIndexTaskClientFactory, + mapper, + emitter, + monitorSchedulerConfig, + rowIngestionMetersFactory + ); + this.awsCredentialsConfig = awsCredentialsConfig; + } + + + @Override + public Supervisor createSupervisor() + { + return new KinesisSupervisor( + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + (KinesisIndexTaskClientFactory) indexTaskClientFactory, + mapper, + this, + rowIngestionMetersFactory, + awsCredentialsConfig + ); + } + + @Override + public String toString() + { + return "KinesisSupervisorSpec{" + + "dataSchema=" + getDataSchema() + + ", tuningConfig=" + getTuningConfig() + + ", ioConfig=" + getIoConfig() + + ", suspended=" + isSuspended() + + ", context=" + getContext() + + '}'; + } + + @Override + @JsonProperty + public KinesisSupervisorTuningConfig getTuningConfig() + { + return (KinesisSupervisorTuningConfig) super.getTuningConfig(); + } + + @Override + @JsonProperty + public KinesisSupervisorIOConfig getIoConfig() + { + return (KinesisSupervisorIOConfig) super.getIoConfig(); + } + + @Override + protected KinesisSupervisorSpec toggleSuspend(boolean suspend) + { + return new KinesisSupervisorSpec( + getDataSchema(), + getTuningConfig(), + getIoConfig(), + getContext(), + suspend, + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + (KinesisIndexTaskClientFactory) indexTaskClientFactory, + mapper, + emitter, + monitorSchedulerConfig, + rowIngestionMetersFactory, + awsCredentialsConfig + ); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfig.java b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfig.java new file mode 100644 index 000000000000..3c749be5cffc --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfig.java @@ -0,0 +1,213 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorTuningConfig; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory; +import org.joda.time.Duration; +import org.joda.time.Period; + +import javax.annotation.Nullable; +import java.io.File; + +public class KinesisSupervisorTuningConfig extends KinesisIndexTaskTuningConfig + implements SeekableStreamSupervisorTuningConfig +{ + private final Integer workerThreads; + private final Integer chatThreads; + private final Long chatRetries; + private final Duration httpTimeout; + private final Duration shutdownTimeout; + + public KinesisSupervisorTuningConfig( + @JsonProperty("maxRowsInMemory") Integer maxRowsInMemory, + @JsonProperty("maxBytesInMemory") Long maxBytesInMemory, + @JsonProperty("maxRowsPerSegment") Integer maxRowsPerSegment, + @JsonProperty("maxTotalRows") Long maxTotalRows, + @JsonProperty("intermediatePersistPeriod") Period intermediatePersistPeriod, + @JsonProperty("basePersistDirectory") File basePersistDirectory, + @JsonProperty("maxPendingPersists") Integer maxPendingPersists, + @JsonProperty("indexSpec") IndexSpec indexSpec, + @JsonProperty("buildV9Directly") Boolean buildV9Directly, + @JsonProperty("reportParseExceptions") Boolean reportParseExceptions, + @JsonProperty("handoffConditionTimeout") Long handoffConditionTimeout, + @JsonProperty("resetOffsetAutomatically") Boolean resetOffsetAutomatically, + @JsonProperty("skipSequenceNumberAvailabilityCheck") Boolean skipSequenceNumberAvailabilityCheck, + @JsonProperty("segmentWriteOutMediumFactory") @Nullable SegmentWriteOutMediumFactory segmentWriteOutMediumFactory, + @JsonProperty("workerThreads") Integer workerThreads, + @JsonProperty("chatThreads") Integer chatThreads, + @JsonProperty("chatRetries") Long chatRetries, + @JsonProperty("httpTimeout") Period httpTimeout, + @JsonProperty("shutdownTimeout") Period shutdownTimeout, + @JsonProperty("recordBufferSize") Integer recordBufferSize, + @JsonProperty("recordBufferOfferTimeout") Integer recordBufferOfferTimeout, + @JsonProperty("recordBufferFullWait") Integer recordBufferFullWait, + @JsonProperty("fetchSequenceNumberTimeout") Integer fetchSequenceNumberTimeout, + @JsonProperty("fetchThreads") Integer fetchThreads, + @JsonProperty("logParseExceptions") @Nullable Boolean logParseExceptions, + @JsonProperty("maxParseExceptions") @Nullable Integer maxParseExceptions, + @JsonProperty("maxSavedParseExceptions") @Nullable Integer maxSavedParseExceptions, + @JsonProperty("maxRecordsPerPoll") @Nullable Integer maxRecordsPerPoll, + @JsonProperty("intermediateHandoffPeriod") Period intermediateHandoffPeriod + ) + { + super( + maxRowsInMemory, + maxBytesInMemory, + maxRowsPerSegment, + maxTotalRows, + intermediatePersistPeriod, + basePersistDirectory, + maxPendingPersists, + indexSpec, + buildV9Directly, + reportParseExceptions, + handoffConditionTimeout, + resetOffsetAutomatically, + skipSequenceNumberAvailabilityCheck, + recordBufferSize, + recordBufferOfferTimeout, + recordBufferFullWait, + fetchSequenceNumberTimeout, + fetchThreads, + segmentWriteOutMediumFactory, + logParseExceptions, + maxParseExceptions, + maxSavedParseExceptions, + maxRecordsPerPoll, + intermediateHandoffPeriod + ); + + this.workerThreads = workerThreads; + this.chatThreads = chatThreads; + this.chatRetries = (chatRetries != null ? chatRetries : DEFAULT_CHAT_RETRIES); + this.httpTimeout = SeekableStreamSupervisorTuningConfig.defaultDuration(httpTimeout, DEFAULT_HTTP_TIMEOUT); + this.shutdownTimeout = SeekableStreamSupervisorTuningConfig.defaultDuration( + shutdownTimeout, + DEFAULT_SHUTDOWN_TIMEOUT + ); + } + + @Override + @JsonProperty + public Integer getWorkerThreads() + { + return workerThreads; + } + + @Override + @JsonProperty + public Integer getChatThreads() + { + return chatThreads; + } + + @Override + @JsonProperty + public Long getChatRetries() + { + return chatRetries; + } + + @Override + @JsonProperty + public Duration getHttpTimeout() + { + return httpTimeout; + } + + @Override + @JsonProperty + public Duration getShutdownTimeout() + { + return shutdownTimeout; + } + + @Override + public String toString() + { + return "KinesisSupervisorTuningConfig{" + + "maxRowsInMemory=" + getMaxRowsInMemory() + + ", maxBytesInMemory=" + getMaxBytesInMemory() + + ", maxRowsPerSegment=" + getMaxRowsPerSegment() + + ", maxTotalRows=" + getMaxTotalRows() + + ", intermediatePersistPeriod=" + getIntermediatePersistPeriod() + + ", basePersistDirectory=" + getBasePersistDirectory() + + ", maxPendingPersists=" + getMaxPendingPersists() + + ", indexSpec=" + getIndexSpec() + + ", reportParseExceptions=" + isReportParseExceptions() + + ", handoffConditionTimeout=" + getHandoffConditionTimeout() + + ", resetOffsetAutomatically=" + isResetOffsetAutomatically() + + ", skipSequenceNumberAvailabilityCheck=" + isSkipSequenceNumberAvailabilityCheck() + + ", workerThreads=" + workerThreads + + ", chatThreads=" + chatThreads + + ", chatRetries=" + chatRetries + + ", httpTimeout=" + httpTimeout + + ", shutdownTimeout=" + shutdownTimeout + + ", recordBufferSize=" + getRecordBufferSize() + + ", recordBufferOfferTimeout=" + getRecordBufferOfferTimeout() + + ", recordBufferFullWait=" + getRecordBufferFullWait() + + ", fetchSequenceNumberTimeout=" + getFetchSequenceNumberTimeout() + + ", fetchThreads=" + getFetchThreads() + + ", segmentWriteOutMediumFactory=" + getSegmentWriteOutMediumFactory() + + ", logParseExceptions=" + isLogParseExceptions() + + ", maxParseExceptions=" + getMaxParseExceptions() + + ", maxSavedParseExceptions=" + getMaxSavedParseExceptions() + + ", maxRecordsPerPoll=" + getMaxRecordsPerPoll() + + ", intermediateHandoffPeriod=" + getIntermediateHandoffPeriod() + + '}'; + } + + @Override + public SeekableStreamIndexTaskTuningConfig convertToTaskTuningConfig() + { + return new KinesisIndexTaskTuningConfig( + getMaxRowsInMemory(), + getMaxBytesInMemory(), + getMaxRowsPerSegment(), + getMaxTotalRows(), + getIntermediatePersistPeriod(), + getBasePersistDirectory(), + getMaxPendingPersists(), + getIndexSpec(), + true, + isReportParseExceptions(), + getHandoffConditionTimeout(), + isResetOffsetAutomatically(), + isSkipSequenceNumberAvailabilityCheck(), + getRecordBufferSize(), + getRecordBufferOfferTimeout(), + getRecordBufferFullWait(), + getFetchSequenceNumberTimeout(), + getFetchThreads(), + getSegmentWriteOutMediumFactory(), + isLogParseExceptions(), + getMaxParseExceptions(), + getMaxSavedParseExceptions(), + getMaxRecordsPerPoll(), + getIntermediateHandoffPeriod() + ); + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/main/resources/META-INF/services/org.apache.druid.initialization.DruidModule b/extensions-core/kinesis-indexing-service/src/main/resources/META-INF/services/org.apache.druid.initialization.DruidModule new file mode 100644 index 000000000000..4c1d6cb0a28e --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/main/resources/META-INF/services/org.apache.druid.initialization.DruidModule @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +org.apache.druid.indexing.kinesis.KinesisIndexingServiceModule diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadataTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadataTest.java new file mode 100644 index 000000000000..f1e3b0fca65e --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisDataSourceMetadataTest.java @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + + +package org.apache.druid.indexing.kinesis; + +import com.google.common.collect.ImmutableMap; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Map; + +public class KinesisDataSourceMetadataTest +{ + private static final KinesisDataSourceMetadata KM0 = KM("foo", ImmutableMap.of()); + private static final KinesisDataSourceMetadata KM1 = KM("foo", ImmutableMap.of("0", "2L", "1", "3L")); + private static final KinesisDataSourceMetadata KM2 = KM("foo", ImmutableMap.of("0", "2L", "1", "4L", "2", "5L")); + private static final KinesisDataSourceMetadata KM3 = KM("foo", ImmutableMap.of("0", "2L", "2", "5L")); + + @Test + public void testMatches() + { + Assert.assertTrue(KM0.matches(KM0)); + Assert.assertTrue(KM0.matches(KM1)); + Assert.assertTrue(KM0.matches(KM2)); + Assert.assertTrue(KM0.matches(KM3)); + + Assert.assertTrue(KM1.matches(KM0)); + Assert.assertTrue(KM1.matches(KM1)); + Assert.assertFalse(KM1.matches(KM2)); + Assert.assertTrue(KM1.matches(KM3)); + + Assert.assertTrue(KM2.matches(KM0)); + Assert.assertFalse(KM2.matches(KM1)); + Assert.assertTrue(KM2.matches(KM2)); + Assert.assertTrue(KM2.matches(KM3)); + + Assert.assertTrue(KM3.matches(KM0)); + Assert.assertTrue(KM3.matches(KM1)); + Assert.assertTrue(KM3.matches(KM2)); + Assert.assertTrue(KM3.matches(KM3)); + } + + @Test + public void testIsValidStart() + { + Assert.assertTrue(KM0.isValidStart()); + Assert.assertTrue(KM1.isValidStart()); + Assert.assertTrue(KM2.isValidStart()); + Assert.assertTrue(KM3.isValidStart()); + } + + @Test + public void testPlus() + { + Assert.assertEquals( + KM("foo", ImmutableMap.of("0", "2L", "1", "3L", "2", "5L")), + KM1.plus(KM3) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of("0", "2L", "1", "4L", "2", "5L")), + KM0.plus(KM2) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of("0", "2L", "1", "4L", "2", "5L")), + KM1.plus(KM2) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of("0", "2L", "1", "3L", "2", "5L")), + KM2.plus(KM1) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of("0", "2L", "1", "4L", "2", "5L")), + KM2.plus(KM2) + ); + } + + @Test + public void testMinus() + { + Assert.assertEquals( + KM("foo", ImmutableMap.of("1", "3L")), + KM1.minus(KM3) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of()), + KM0.minus(KM2) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of()), + KM1.minus(KM2) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of("2", "5L")), + KM2.minus(KM1) + ); + + Assert.assertEquals( + KM("foo", ImmutableMap.of()), + KM2.minus(KM2) + ); + } + + private static KinesisDataSourceMetadata KM(String stream, Map sequences) + { + return new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, sequences)); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIOConfigTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIOConfigTest.java new file mode 100644 index 000000000000..fff34f946b26 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIOConfigTest.java @@ -0,0 +1,249 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.segment.indexing.IOConfig; +import org.hamcrest.CoreMatchers; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +import java.util.Collections; + +public class KinesisIOConfigTest +{ + private final ObjectMapper mapper; + + public KinesisIOConfigTest() + { + mapper = new DefaultObjectMapper(); + mapper.registerModules((Iterable) new KinesisIndexingServiceModule().getJacksonModules()); + } + + @Rule + public final ExpectedException exception = ExpectedException.none(); + + @Test + public void testSerdeWithDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}}\n" + + "}"; + + KinesisIndexTaskIOConfig config = (KinesisIndexTaskIOConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + IOConfig.class + ) + ), IOConfig.class + ); + + Assert.assertNull(config.getTaskGroupId()); + Assert.assertEquals("my-sequence-name", config.getBaseSequenceName()); + Assert.assertEquals("mystream", config.getStartPartitions().getStream()); + Assert.assertEquals( + ImmutableMap.of("0", "1", "1", "10"), + config.getStartPartitions().getPartitionSequenceNumberMap() + ); + Assert.assertEquals("mystream", config.getEndPartitions().getStream()); + Assert.assertEquals( + ImmutableMap.of("0", "15", "1", "200"), + config.getEndPartitions().getPartitionSequenceNumberMap() + ); + Assert.assertTrue(config.isUseTransaction()); + Assert.assertFalse("minimumMessageTime", config.getMinimumMessageTime().isPresent()); + Assert.assertEquals(config.getEndpoint(), "kinesis.us-east-1.amazonaws.com"); + Assert.assertEquals(config.getRecordsPerFetch(), 4000); + Assert.assertEquals(config.getFetchDelayMillis(), 0); + Assert.assertEquals(Collections.emptySet(), config.getExclusiveStartSequenceNumberPartitions()); + Assert.assertNull(config.getAwsAssumedRoleArn()); + Assert.assertNull(config.getAwsExternalId()); + Assert.assertFalse(config.isDeaggregate()); + Assert.assertTrue(config.isSkipOffsetGaps()); + } + + @Test + public void testSerdeWithNonDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"taskGroupId\": 0,\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}},\n" + + " \"useTransaction\": false,\n" + + " \"minimumMessageTime\": \"2016-05-31T12:00Z\",\n" + + " \"maximumMessageTime\": \"2016-05-31T14:00Z\",\n" + + " \"endpoint\": \"kinesis.us-east-2.amazonaws.com\",\n" + + " \"recordsPerFetch\": 1000,\n" + + " \"fetchDelayMillis\": 1000,\n" + + " \"exclusiveStartSequenceNumberPartitions\": [\"0\"],\n" + + " \"awsAssumedRoleArn\": \"role\",\n" + + " \"awsExternalId\": \"awsexternalid\",\n" + + " \"deaggregate\": true\n" + + "}"; + + KinesisIndexTaskIOConfig config = (KinesisIndexTaskIOConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + IOConfig.class + ) + ), IOConfig.class + ); + + Assert.assertEquals((Integer) 0, config.getTaskGroupId()); + Assert.assertEquals("my-sequence-name", config.getBaseSequenceName()); + Assert.assertEquals("mystream", config.getStartPartitions().getStream()); + Assert.assertEquals( + ImmutableMap.of("0", "1", "1", "10"), + config.getStartPartitions().getPartitionSequenceNumberMap() + ); + Assert.assertEquals("mystream", config.getEndPartitions().getStream()); + Assert.assertEquals( + ImmutableMap.of("0", "15", "1", "200"), + config.getEndPartitions().getPartitionSequenceNumberMap() + ); + Assert.assertFalse(config.isUseTransaction()); + Assert.assertTrue("maximumMessageTime", config.getMaximumMessageTime().isPresent()); + Assert.assertTrue("minimumMessageTime", config.getMinimumMessageTime().isPresent()); + Assert.assertEquals(DateTimes.of("2016-05-31T12:00Z"), config.getMinimumMessageTime().get()); + Assert.assertEquals(DateTimes.of("2016-05-31T14:00Z"), config.getMaximumMessageTime().get()); + Assert.assertEquals(config.getEndpoint(), "kinesis.us-east-2.amazonaws.com"); + Assert.assertEquals(config.getExclusiveStartSequenceNumberPartitions(), ImmutableSet.of("0")); + Assert.assertEquals(1000, config.getRecordsPerFetch()); + Assert.assertEquals(1000, config.getFetchDelayMillis()); + Assert.assertEquals("role", config.getAwsAssumedRoleArn()); + Assert.assertEquals("awsexternalid", config.getAwsExternalId()); + Assert.assertTrue(config.isDeaggregate()); + Assert.assertTrue(config.isSkipOffsetGaps()); + } + + @Test + public void testBaseSequenceNameRequired() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(NullPointerException.class)); + exception.expectMessage(CoreMatchers.containsString("baseSequenceName")); + mapper.readValue(jsonStr, IOConfig.class); + } + + @Test + public void testStartPartitionsRequired() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(NullPointerException.class)); + exception.expectMessage(CoreMatchers.containsString("startPartitions")); + mapper.readValue(jsonStr, IOConfig.class); + } + + @Test + public void testEndPartitionsRequired() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(NullPointerException.class)); + exception.expectMessage(CoreMatchers.containsString("endPartitions")); + mapper.readValue(jsonStr, IOConfig.class); + } + + @Test + public void testStartAndEndstreamMatch() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"notmystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(IllegalArgumentException.class)); + exception.expectMessage(CoreMatchers.containsString("must match")); + mapper.readValue(jsonStr, IOConfig.class); + } + + @Test + public void testStartAndEndPartitionSetMatch() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"endpoint\": \"kinesis.us-east-1.amazonaws.com\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"2\":\"200\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(IllegalArgumentException.class)); + exception.expectMessage(CoreMatchers.containsString("start partition set and end partition set must match")); + mapper.readValue(jsonStr, IOConfig.class); + } + + @Test + public void testEndPointRequired() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"baseSequenceName\": \"my-sequence-name\",\n" + + " \"startPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"1\", \"1\":\"10\"}},\n" + + " \"endPartitions\": {\"stream\":\"mystream\", \"partitionSequenceNumberMap\" : {\"0\":\"15\", \"1\":\"200\"}}\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(NullPointerException.class)); + exception.expectMessage(CoreMatchers.containsString("endpoint")); + mapper.readValue(jsonStr, IOConfig.class); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientTest.java new file mode 100644 index 000000000000..9d7fafbb4c73 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientTest.java @@ -0,0 +1,1059 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Optional; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Maps; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.indexer.TaskLocation; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.IndexTaskClient; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.IAE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.java.util.http.client.Request; +import org.apache.druid.java.util.http.client.response.FullResponseHandler; +import org.apache.druid.java.util.http.client.response.FullResponseHolder; +import org.easymock.Capture; +import org.easymock.CaptureType; +import org.easymock.EasyMock; +import org.easymock.EasyMockSupport; +import org.jboss.netty.handler.codec.http.HttpHeaders; +import org.jboss.netty.handler.codec.http.HttpMethod; +import org.jboss.netty.handler.codec.http.HttpResponse; +import org.jboss.netty.handler.codec.http.HttpResponseStatus; +import org.joda.time.DateTime; +import org.joda.time.Duration; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.io.IOException; +import java.net.URL; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.easymock.EasyMock.expect; + + +@RunWith(Parameterized.class) +public class KinesisIndexTaskClientTest extends EasyMockSupport +{ + @Rule + public ExpectedException expectedException = ExpectedException.none(); + + private static final ObjectMapper objectMapper = new DefaultObjectMapper(); + private static final String TEST_ID = "test-id"; + private static final List TEST_IDS = Arrays.asList("test-id1", "test-id2", "test-id3", "test-id4"); + private static final String TEST_HOST = "test-host"; + private static final int TEST_PORT = 1234; + private static final int TEST_TLS_PORT = -1; + private static final String TEST_DATASOURCE = "test-datasource"; + private static final Duration TEST_HTTP_TIMEOUT = new Duration(5000); + private static final long TEST_NUM_RETRIES = 0; + private static final String URL_FORMATTER = "http://%s:%d/druid/worker/v1/chat/%s/%s"; + + private int numThreads; + private HttpClient httpClient; + private TaskInfoProvider taskInfoProvider; + private FullResponseHolder responseHolder; + private HttpResponse response; + private HttpHeaders headers; + private KinesisIndexTaskClient client; + + @Parameterized.Parameters(name = "numThreads = {0}") + public static Iterable constructorFeeder() + { + return ImmutableList.of(new Object[]{1}, new Object[]{8}); + } + + public KinesisIndexTaskClientTest(int numThreads) + { + this.numThreads = numThreads; + } + + @Before + public void setUp() + { + httpClient = createMock(HttpClient.class); + taskInfoProvider = createMock(TaskInfoProvider.class); + responseHolder = createMock(FullResponseHolder.class); + response = createMock(HttpResponse.class); + headers = createMock(HttpHeaders.class); + + client = new TestableKinesisIndexTaskClient(httpClient, objectMapper, taskInfoProvider); + expect(taskInfoProvider.getTaskLocation(TEST_ID)) + .andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) + .anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.running(TEST_ID))) + .anyTimes(); + + for (String testId : TEST_IDS) { + expect(taskInfoProvider.getTaskLocation(testId)) + .andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) + .anyTimes(); + expect(taskInfoProvider.getTaskStatus(testId)) + .andReturn(Optional.of(TaskStatus.running(testId))) + .anyTimes(); + } + } + + @After + public void tearDown() + { + client.close(); + } + + @Test + public void testNoTaskLocation() throws IOException + { + EasyMock.reset(taskInfoProvider); + expect(taskInfoProvider.getTaskLocation(TEST_ID)).andReturn(TaskLocation.unknown()).anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.running(TEST_ID))) + .anyTimes(); + replayAll(); + + Assert.assertFalse(client.stop(TEST_ID, true)); + Assert.assertFalse(client.resume(TEST_ID)); + Assert.assertEquals(ImmutableMap.of(), client.pause(TEST_ID)); + Assert.assertEquals(ImmutableMap.of(), client.pause(TEST_ID)); + Assert.assertEquals(Status.NOT_STARTED, client.getStatus(TEST_ID)); + Assert.assertNull(client.getStartTime(TEST_ID)); + Assert.assertEquals(ImmutableMap.of(), client.getCurrentOffsets(TEST_ID, true)); + Assert.assertEquals(ImmutableMap.of(), client.getEndOffsets(TEST_ID)); + Assert.assertFalse(client.setEndOffsets(TEST_ID, Collections.emptyMap(), true)); + Assert.assertFalse(client.setEndOffsets(TEST_ID, Collections.emptyMap(), true)); + + verifyAll(); + } + + @Test + public void testTaskNotRunnableException() + { + expectedException.expect(IndexTaskClient.TaskNotRunnableException.class); + expectedException.expectMessage("Aborting request because task [test-id] is not runnable"); + + EasyMock.reset(taskInfoProvider); + expect(taskInfoProvider.getTaskLocation(TEST_ID)) + .andReturn(new TaskLocation(TEST_HOST, TEST_PORT, TEST_TLS_PORT)) + .anyTimes(); + expect(taskInfoProvider.getTaskStatus(TEST_ID)) + .andReturn(Optional.of(TaskStatus.failure(TEST_ID))) + .anyTimes(); + replayAll(); + + client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + } + + @Test + public void testInternalServerError() + { + expectedException.expect(RuntimeException.class); + expectedException.expectMessage("org.apache.druid.java.util.common.IOE: Received status [500]"); + + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.INTERNAL_SERVER_ERROR).times(2); + expect( + httpClient.go( + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + ) + ).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + } + + @Test + public void testBadRequest() + { + expectedException.expect(IAE.class); + expectedException.expectMessage("Received 400 Bad Request with body:"); + + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.BAD_REQUEST).times(2); + expect(responseHolder.getContent()).andReturn(""); + expect( + httpClient.go( + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + ) + ).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + } + + @Test + public void testTaskLocationMismatch() + { + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.NOT_FOUND).times(3) + .andReturn(HttpResponseStatus.OK); + expect(responseHolder.getResponse()).andReturn(response); + expect(responseHolder.getContent()).andReturn("") + .andReturn("{}"); + expect(response.headers()).andReturn(headers); + expect(headers.get("X-Druid-Task-Id")).andReturn("a-different-task-id"); + expect( + httpClient.go( + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + ) + ).andReturn( + Futures.immediateFuture(responseHolder) + ).times(2); + replayAll(); + + Map results = client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + + Assert.assertEquals(0, results.size()); + } + + @Test + public void testGetCurrentOffsets() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); + expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}"); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + Map results = client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/current"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + Assert.assertEquals(2, results.size()); + Assert.assertEquals("1", results.get("0")); + Assert.assertEquals("10", results.get("1")); + } + + @Test + public void testGetCurrentOffsetsWithRetry() throws Exception + { + client = new TestableKinesisIndexTaskClient(httpClient, objectMapper, taskInfoProvider, 3); + + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.NOT_FOUND).times(6) + .andReturn(HttpResponseStatus.OK).times(1); + expect(responseHolder.getContent()).andReturn("").times(2) + .andReturn("{\"0\":1, \"1\":10}"); + expect(responseHolder.getResponse()).andReturn(response).times(2); + expect(response.headers()).andReturn(headers).times(2); + expect(headers.get("X-Druid-Task-Id")).andReturn(TEST_ID).times(2); + + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(3); + + replayAll(); + + Map results = client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + + Assert.assertEquals(3, captured.getValues().size()); + for (Request request : captured.getValues()) { + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/current"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + } + + Assert.assertEquals(2, results.size()); + Assert.assertEquals("1", results.get("0")); + Assert.assertEquals("10", results.get("1")); + } + + @Test + public void testGetCurrentOffsetsWithExhaustedRetries() + { + expectedException.expect(RuntimeException.class); + expectedException.expectMessage("org.apache.druid.java.util.common.IOE: Received status [404]"); + + client = new TestableKinesisIndexTaskClient(httpClient, objectMapper, taskInfoProvider, 2); + + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.NOT_FOUND).anyTimes(); + expect(responseHolder.getContent()).andReturn("").anyTimes(); + expect(responseHolder.getResponse()).andReturn(response).anyTimes(); + expect(response.headers()).andReturn(headers).anyTimes(); + expect(headers.get("X-Druid-Task-Id")).andReturn(TEST_ID).anyTimes(); + + expect( + httpClient.go( + EasyMock.anyObject(Request.class), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + ) + ).andReturn(Futures.immediateFuture(responseHolder)).anyTimes(); + replayAll(); + + client.getCurrentOffsets(TEST_ID, true); + verifyAll(); + } + + @Test + public void testGetEndOffsets() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); + expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}"); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + Map results = client.getEndOffsets(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/end"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + Assert.assertEquals(2, results.size()); + Assert.assertEquals("1", results.get("0")); + Assert.assertEquals("10", results.get("1")); + } + + @Test + public void testGetStartTime() throws Exception + { + client = new TestableKinesisIndexTaskClient(httpClient, objectMapper, taskInfoProvider, 2); + DateTime now = DateTimes.nowUtc(); + + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.NOT_FOUND).times(3) + .andReturn(HttpResponseStatus.OK); + expect(responseHolder.getResponse()).andReturn(response); + expect(response.headers()).andReturn(headers); + expect(headers.get("X-Druid-Task-Id")).andReturn(null); + expect(responseHolder.getContent()).andReturn(String.valueOf(now.getMillis())).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(2); + replayAll(); + + DateTime results = client.getStartTime(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/time/start"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + Assert.assertEquals(now, results); + } + + @Test + public void testGetStatus() throws Exception + { + Status status = Status.READING; + + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK); + expect(responseHolder.getContent()).andReturn(StringUtils.format("\"%s\"", status.toString())).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + Status results = client.getStatus(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/status"), + request.getUrl() + ); + Assert.assertTrue(null, request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + Assert.assertEquals(status, results); + } + + @Test + public void testPause() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).times(2); + expect(responseHolder.getContent()).andReturn("{\"0\":1, \"1\":10}").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + Map results = client.pause(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/pause"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + Assert.assertEquals(2, results.size()); + Assert.assertEquals("1", results.get("0")); + Assert.assertEquals("10", results.get("1")); + } + + @Test + public void testPauseWithSubsequentGetOffsets() throws Exception + { + Capture captured = Capture.newInstance(); + Capture captured2 = Capture.newInstance(); + Capture captured3 = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.ACCEPTED).times(2) + .andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn("\"PAUSED\"").times(2) + .andReturn("{\"0\":1, \"1\":10}").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + expect(httpClient.go( + EasyMock.capture(captured2), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + expect(httpClient.go( + EasyMock.capture(captured3), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + + replayAll(); + + Map results = client.pause(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/pause"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + + request = captured2.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/status"), + request.getUrl() + ); + + request = captured3.getValue(); + Assert.assertEquals(HttpMethod.GET, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/current"), + request.getUrl() + ); + + Assert.assertEquals(2, results.size()); + Assert.assertEquals("1", results.get("0")); + Assert.assertEquals("10", results.get("1")); + } + + @Test + public void testResume() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.resume(TEST_ID); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/resume"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + } + + @Test + public void testSetEndOffsets() throws Exception + { + Map endOffsets = ImmutableMap.of("0", "15", "1", "120"); + + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.setEndOffsets(TEST_ID, endOffsets, true); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/end?finish=true"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + Assert.assertEquals("{\"0\":\"15\",\"1\":\"120\"}", StringUtils.fromUtf8(request.getContent().array())); + } + + @Test + public void testSetEndOffsetsAndResume() throws Exception + { + Map endOffsets = ImmutableMap.of("0", "15", "1", "120"); + + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.setEndOffsets(TEST_ID, endOffsets, true); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/offsets/end?finish=true"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + Assert.assertEquals("{\"0\":\"15\",\"1\":\"120\"}", StringUtils.fromUtf8(request.getContent().array())); + } + + @Test + public void testStop() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.stop(TEST_ID, false); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/stop"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + } + + @Test + public void testStopAndPublish() throws Exception + { + Capture captured = Capture.newInstance(); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ); + replayAll(); + + client.stop(TEST_ID, true); + verifyAll(); + + Request request = captured.getValue(); + Assert.assertEquals(HttpMethod.POST, request.getMethod()); + Assert.assertEquals( + new URL("http://test-host:1234/druid/worker/v1/chat/test-id/stop?publish=true"), + request.getUrl() + ); + Assert.assertTrue(request.getHeaders().get("X-Druid-Task-Id").contains("test-id")); + } + + @Test + public void testStopAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "stop"))); + futures.add(client.stopAsync(testId, false)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.POST, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertTrue(responses.get(i)); + } + } + + @Test + public void testResumeAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "resume"))); + futures.add(client.resumeAsync(testId)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.POST, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertTrue(responses.get(i)); + } + } + + @Test + public void testPauseAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List>> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "pause"))); + futures.add(client.pauseAsync(testId)); + } + + List> responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.POST, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertEquals(Maps.newLinkedHashMap(ImmutableMap.of("0", "1")), responses.get(i)); + } + } + + @Test + public void testGetStatusAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn("\"READING\"").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "status"))); + futures.add(client.getStatusAsync(testId)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.GET, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertEquals(Status.READING, responses.get(i)); + } + } + + @Test + public void testGetStartTimeAsync() throws Exception + { + final DateTime now = DateTimes.nowUtc(); + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn(String.valueOf(now.getMillis())).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "time/start"))); + futures.add(client.getStartTimeAsync(testId)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.GET, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertEquals(now, responses.get(i)); + } + } + + @Test + public void testGetCurrentOffsetsAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List>> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "offsets/current"))); + futures.add(client.getCurrentOffsetsAsync(testId, false)); + } + + List> responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.GET, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertEquals(Maps.newLinkedHashMap(ImmutableMap.of("0", "1")), responses.get(i)); + } + } + + @Test + public void testGetEndOffsetsAsync() throws Exception + { + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(responseHolder.getContent()).andReturn("{\"0\":\"1\"}").anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List>> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format(URL_FORMATTER, TEST_HOST, TEST_PORT, testId, "offsets/end"))); + futures.add(client.getEndOffsetsAsync(testId)); + } + + List> responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.GET, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertEquals(Maps.newLinkedHashMap(ImmutableMap.of("0", "1")), responses.get(i)); + } + } + + @Test + public void testSetEndOffsetsAsync() throws Exception + { + Map endOffsets = ImmutableMap.of("0", "15L", "1", "120L"); + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add(new URL(StringUtils.format( + URL_FORMATTER, + TEST_HOST, + TEST_PORT, + testId, + StringUtils.format("offsets/end?finish=%s", true) + ))); + futures.add(client.setEndOffsetsAsync(testId, endOffsets, true)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.POST, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertTrue(responses.get(i)); + } + } + + @Test + public void testSetEndOffsetsAsyncWithResume() throws Exception + { + Map endOffsets = ImmutableMap.of("0", "15L", "1", "120L"); + final int numRequests = TEST_IDS.size(); + Capture captured = Capture.newInstance(CaptureType.ALL); + expect(responseHolder.getStatus()).andReturn(HttpResponseStatus.OK).anyTimes(); + expect(httpClient.go( + EasyMock.capture(captured), + EasyMock.anyObject(FullResponseHandler.class), + EasyMock.eq(TEST_HTTP_TIMEOUT) + )).andReturn( + Futures.immediateFuture(responseHolder) + ).times(numRequests); + replayAll(); + + List expectedUrls = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (String testId : TEST_IDS) { + expectedUrls.add( + new URL( + StringUtils.format( + URL_FORMATTER, + TEST_HOST, + TEST_PORT, + testId, + "offsets/end?finish=true" + ) + ) + ); + futures.add(client.setEndOffsetsAsync(testId, endOffsets, true)); + } + + List responses = Futures.allAsList(futures).get(); + + verifyAll(); + List requests = captured.getValues(); + + Assert.assertEquals(numRequests, requests.size()); + Assert.assertEquals(numRequests, responses.size()); + for (int i = 0; i < numRequests; i++) { + Assert.assertEquals(HttpMethod.POST, requests.get(i).getMethod()); + Assert.assertTrue("unexpectedURL", expectedUrls.contains(requests.get(i).getUrl())); + Assert.assertTrue(responses.get(i)); + } + } + + private class TestableKinesisIndexTaskClient extends KinesisIndexTaskClient + { + TestableKinesisIndexTaskClient( + HttpClient httpClient, + ObjectMapper jsonMapper, + TaskInfoProvider taskInfoProvider + ) + { + this(httpClient, jsonMapper, taskInfoProvider, TEST_NUM_RETRIES); + } + + TestableKinesisIndexTaskClient( + HttpClient httpClient, + ObjectMapper jsonMapper, + TaskInfoProvider taskInfoProvider, + long numRetries + ) + { + super(httpClient, jsonMapper, taskInfoProvider, TEST_DATASOURCE, numThreads, TEST_HTTP_TIMEOUT, numRetries); + } + + @Override + protected void checkConnection(String host, int port) + { + } + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java new file mode 100644 index 000000000000..cef97955db2e --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java @@ -0,0 +1,2968 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.annotation.JacksonInject; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.InjectableValues; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Predicate; +import com.google.common.base.Predicates; +import com.google.common.base.Throwables; +import com.google.common.collect.FluentIterable; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; +import com.google.common.io.Files; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.MoreExecutors; +import org.apache.druid.client.cache.CacheConfig; +import org.apache.druid.client.cache.CachePopulatorStats; +import org.apache.druid.client.cache.MapCache; +import org.apache.druid.common.aws.AWSCredentialsConfig; +import org.apache.druid.data.input.impl.DimensionsSpec; +import org.apache.druid.data.input.impl.FloatDimensionSchema; +import org.apache.druid.data.input.impl.JSONParseSpec; +import org.apache.druid.data.input.impl.LongDimensionSchema; +import org.apache.druid.data.input.impl.StringDimensionSchema; +import org.apache.druid.data.input.impl.StringInputRowParser; +import org.apache.druid.data.input.impl.TimestampSpec; +import org.apache.druid.discovery.DataNodeService; +import org.apache.druid.discovery.DruidNodeAnnouncer; +import org.apache.druid.discovery.LookupNodeService; +import org.apache.druid.indexer.TaskState; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.Counters; +import org.apache.druid.indexing.common.IngestionStatsAndErrorsTaskReportData; +import org.apache.druid.indexing.common.SegmentLoaderFactory; +import org.apache.druid.indexing.common.TaskLock; +import org.apache.druid.indexing.common.TaskReport; +import org.apache.druid.indexing.common.TaskReportFileWriter; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.TaskToolboxFactory; +import org.apache.druid.indexing.common.TestUtils; +import org.apache.druid.indexing.common.actions.LocalTaskActionClientFactory; +import org.apache.druid.indexing.common.actions.TaskActionClientFactory; +import org.apache.druid.indexing.common.actions.TaskActionToolbox; +import org.apache.druid.indexing.common.actions.TaskAuditLogConfig; +import org.apache.druid.indexing.common.config.TaskConfig; +import org.apache.druid.indexing.common.config.TaskStorageConfig; +import org.apache.druid.indexing.common.stats.RowIngestionMeters; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.IndexTaskTest; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.common.task.TaskResource; +import org.apache.druid.indexing.overlord.DataSourceMetadata; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.MetadataTaskStorage; +import org.apache.druid.indexing.overlord.TaskLockbox; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.supervisor.SupervisorManager; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor; +import org.apache.druid.indexing.test.TestDataSegmentAnnouncer; +import org.apache.druid.indexing.test.TestDataSegmentKiller; +import org.apache.druid.java.util.common.CompressionUtils; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.common.granularity.Granularities; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.java.util.common.parsers.JSONPathSpec; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.java.util.emitter.core.NoopEmitter; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.java.util.metrics.MonitorScheduler; +import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.metadata.DerbyMetadataStorageActionHandlerFactory; +import org.apache.druid.metadata.EntryExistsException; +import org.apache.druid.metadata.IndexerSQLMetadataStorageCoordinator; +import org.apache.druid.metadata.TestDerbyConnector; +import org.apache.druid.query.DefaultQueryRunnerFactoryConglomerate; +import org.apache.druid.query.Druids; +import org.apache.druid.query.IntervalChunkingQueryRunnerDecorator; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.QueryRunnerFactoryConglomerate; +import org.apache.druid.query.QueryToolChest; +import org.apache.druid.query.Result; +import org.apache.druid.query.SegmentDescriptor; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.query.aggregation.DoubleSumAggregatorFactory; +import org.apache.druid.query.aggregation.LongSumAggregatorFactory; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.query.timeseries.TimeseriesQueryEngine; +import org.apache.druid.query.timeseries.TimeseriesQueryQueryToolChest; +import org.apache.druid.query.timeseries.TimeseriesQueryRunnerFactory; +import org.apache.druid.query.timeseries.TimeseriesResultValue; +import org.apache.druid.segment.DimensionHandlerUtils; +import org.apache.druid.segment.IndexIO; +import org.apache.druid.segment.QueryableIndex; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.column.DictionaryEncodedColumn; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.indexing.granularity.UniformGranularitySpec; +import org.apache.druid.segment.loading.DataSegmentPusher; +import org.apache.druid.segment.loading.LocalDataSegmentPusher; +import org.apache.druid.segment.loading.LocalDataSegmentPusherConfig; +import org.apache.druid.segment.loading.SegmentLoaderConfig; +import org.apache.druid.segment.loading.SegmentLoaderLocalCacheManager; +import org.apache.druid.segment.loading.StorageLocationConfig; +import org.apache.druid.segment.realtime.appenderator.AppenderatorImpl; +import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; +import org.apache.druid.segment.realtime.plumber.SegmentHandoffNotifier; +import org.apache.druid.segment.realtime.plumber.SegmentHandoffNotifierFactory; +import org.apache.druid.segment.transform.ExpressionTransform; +import org.apache.druid.segment.transform.TransformSpec; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.coordination.DataSegmentServerAnnouncer; +import org.apache.druid.server.coordination.ServerType; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.timeline.DataSegment; +import org.easymock.EasyMock; +import org.easymock.EasyMockSupport; +import org.joda.time.Interval; +import org.joda.time.Period; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import javax.annotation.Nullable; +import java.io.File; +import java.io.IOException; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +import static org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; +import static org.easymock.EasyMock.anyLong; +import static org.easymock.EasyMock.anyObject; +import static org.easymock.EasyMock.anyString; +import static org.easymock.EasyMock.expect; +import static org.easymock.EasyMock.expectLastCall; +import static org.easymock.EasyMock.reset; + +public class KinesisIndexTaskTest extends EasyMockSupport +{ + private static final Logger log = new Logger(KinesisIndexTaskTest.class); + private static final ObjectMapper objectMapper = TestHelper.makeJsonMapper(); + private static String stream = "stream"; + private static String shardId1 = "1"; + private static String shardId0 = "0"; + private static KinesisRecordSupplier recordSupplier; + private static List> records = ImmutableList.of( + new OrderedPartitionableRecord<>(stream, "1", "0", JB("2008", "a", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "1", JB("2009", "b", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "2", JB("2010", "c", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "3", JB("2011", "d", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "4", JB("2011", "e", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>( + stream, + "1", + "5", + JB("246140482-04-24T15:36:27.903Z", "x", "z", "10", "20.0", "1.0") + ), + new OrderedPartitionableRecord<>(stream, "1", "6", Collections.singletonList(StringUtils.toUtf8("unparseable"))), + new OrderedPartitionableRecord<>(stream, "1", "7", Collections.singletonList(StringUtils.toUtf8("unparseable2"))), + new OrderedPartitionableRecord<>(stream, "1", "8", Collections.singletonList(StringUtils.toUtf8("{}"))), + new OrderedPartitionableRecord<>(stream, "1", "9", JB("2013", "f", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "10", JB("2049", "f", "y", "notanumber", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "11", JB("2049", "f", "y", "10", "notanumber", "1.0")), + new OrderedPartitionableRecord<>(stream, "1", "12", JB("2049", "f", "y", "10", "20.0", "notanumber")), + new OrderedPartitionableRecord<>(stream, "0", "0", JB("2012", "g", "y", "10", "20.0", "1.0")), + new OrderedPartitionableRecord<>(stream, "0", "1", JB("2011", "h", "y", "10", "20.0", "1.0")) + ); + + private static ServiceEmitter emitter; + private static ListeningExecutorService taskExec; + + private final List runningTasks = new ArrayList<>(); + + private long handoffConditionTimeout = 0; + private boolean reportParseExceptions = false; + private boolean logParseExceptions = true; + private Integer maxParseExceptions = null; + private Integer maxSavedParseExceptions = null; + private boolean resetOffsetAutomatically = false; + private boolean doHandoff = true; + private int maxRowsInMemory = 1000; + private Integer maxRowsPerSegment = null; + private Long maxTotalRows = null; + private Period intermediateHandoffPeriod = null; + private int maxRecordsPerPoll; + private boolean skipAvailabilityCheck = false; + + private TaskToolboxFactory toolboxFactory; + private IndexerMetadataStorageCoordinator metadataStorageCoordinator; + private TaskStorage taskStorage; + private TaskLockbox taskLockbox; + private File directory; + private final Set checkpointRequestsHash = new HashSet<>(); + private File reportsFile; + private RowIngestionMetersFactory rowIngestionMetersFactory; + + private static final DataSchema DATA_SCHEMA = new DataSchema( + "test_ds", + objectMapper.convertValue( + new StringInputRowParser( + new JSONParseSpec( + new TimestampSpec("timestamp", "iso", null), + new DimensionsSpec( + Arrays.asList( + new StringDimensionSchema("dim1"), + new StringDimensionSchema("dim1t"), + new StringDimensionSchema("dim2"), + new LongDimensionSchema("dimLong"), + new FloatDimensionSchema("dimFloat") + ), + null, + null + ), + new JSONPathSpec(true, ImmutableList.of()), + ImmutableMap.of() + ), + StandardCharsets.UTF_8.name() + ), + Map.class + ), + new AggregatorFactory[]{ + new DoubleSumAggregatorFactory("met1sum", "met1"), + new CountAggregatorFactory("rows") + }, + new UniformGranularitySpec(Granularities.DAY, Granularities.NONE, null), + null, + objectMapper + ); + + @Rule + public final TemporaryFolder tempFolder = new TemporaryFolder(); + + @Rule + public final TestDerbyConnector.DerbyConnectorRule derby = new TestDerbyConnector.DerbyConnectorRule(); + + @BeforeClass + public static void setupClass() + { + emitter = new ServiceEmitter( + "service", + "host", + new NoopEmitter() + ); + emitter.start(); + EmittingLogger.registerEmitter(emitter); + taskExec = MoreExecutors.listeningDecorator( + Executors.newCachedThreadPool( + Execs.makeThreadFactory("kinesis-task-test-%d") + ) + ); + } + + @Before + public void setupTest() throws IOException, InterruptedException + { + handoffConditionTimeout = 0; + reportParseExceptions = false; + logParseExceptions = true; + maxParseExceptions = null; + maxSavedParseExceptions = null; + skipAvailabilityCheck = false; + doHandoff = true; + reportsFile = File.createTempFile("KinesisIndexTaskTestReports-" + System.currentTimeMillis(), "json"); + maxRecordsPerPoll = 1; + + recordSupplier = mock(KinesisRecordSupplier.class); + + // sleep required because of kinesalite + Thread.sleep(500); + makeToolboxFactory(); + } + + @After + public void tearDownTest() + { + synchronized (runningTasks) { + for (Task task : runningTasks) { + task.stopGracefully(); + } + + runningTasks.clear(); + } + reportsFile.delete(); + destroyToolboxFactory(); + } + + @AfterClass + public static void tearDownClass() throws Exception + { + taskExec.shutdown(); + taskExec.awaitTermination(9999, TimeUnit.DAYS); + emitter.close(); + } + + @Test(timeout = 120_000L) + public void testRunAfterDataInserted() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 5)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + @Test(timeout = 120_000L) + public void testRunBeforeDataInserted() throws Exception + { + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(Collections.emptyList()) + .times(5) + .andReturn(records.subList(13, 15)) + .once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + // Check metrics + Assert.assertEquals(2, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2011/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2012/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId0, + "1" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("h"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("g"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testIncrementalHandOff() throws Exception + { + final String baseSequenceName = "sequence0"; + // as soon as any segment has more than one record, incremental publishing should happen + maxRowsPerSegment = 2; + maxRecordsPerPoll = 1; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(0, 5)) + .once() + .andReturn(records.subList(4, records.size())) + .once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + ) + ); + + final SeekableStreamPartitions checkpoint1 = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4", + shardId0, + "0" + ) + ); + + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "9", + shardId0, + "1" + ) + ); + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + baseSequenceName, + startPartitions, + endPartitions, + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + final ListenableFuture future = runTask(task); + while (task.getRunner().getStatus() != SeekableStreamIndexTaskRunner.Status.PAUSED) { + Thread.sleep(10); + } + final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); + Assert.assertTrue(checkpoint1.getPartitionSequenceNumberMap().equals(currentOffsets)); + task.getRunner().setEndOffsets(currentOffsets, false); + + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + Assert.assertEquals(1, checkpointRequestsHash.size()); + Assert.assertTrue( + checkpointRequestsHash.contains( + Objects.hash( + DATA_SCHEMA.getDataSource(), + 0, + new KinesisDataSourceMetadata(startPartitions), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, currentOffsets)) + ) + ) + ); + + // Check metrics + Assert.assertEquals(8, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(4, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2008/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2009/P1D", 0); + SegmentDescriptor desc3 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc4 = SD(task, "2011/P1D", 0); + SegmentDescriptor desc5 = SD(task, "2011/P1D", 1); + SegmentDescriptor desc6 = SD(task, "2012/P1D", 0); + SegmentDescriptor desc7 = SD(task, "2013/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4, desc5, desc6, desc7), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "9", + shardId0, + "1" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("a"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("b"), readSegmentColumn("dim1", desc2)); + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc3)); + Assert.assertTrue((ImmutableList.of("d", "e").equals(readSegmentColumn("dim1", desc4)) + && ImmutableList.of("h").equals(readSegmentColumn("dim1", desc5))) || + (ImmutableList.of("d", "h").equals(readSegmentColumn("dim1", desc4)) + && ImmutableList.of("e").equals(readSegmentColumn("dim1", desc5)))); + Assert.assertEquals(ImmutableList.of("g"), readSegmentColumn("dim1", desc6)); + Assert.assertEquals(ImmutableList.of("f"), readSegmentColumn("dim1", desc7)); + } + + + @Test(timeout = 120_000L) + public void testIncrementalHandOffMaxTotalRows() throws Exception + { + final String baseSequenceName = "sequence0"; + // incremental publish should happen every 3 records + maxRowsPerSegment = Integer.MAX_VALUE; + maxTotalRows = 3L; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(0, 3)) + .once() + .andReturn(records.subList(2, 10)) + .once() + .andReturn(records.subList(9, 11)); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + // Insert data + final SeekableStreamPartitions startPartitions = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "0" + ) + ); + // Checkpointing will happen at either checkpoint1 or checkpoint2 depending on ordering + // of events fetched across two partitions from Kafka + final SeekableStreamPartitions checkpoint1 = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "2" + ) + ); + final SeekableStreamPartitions checkpoint2 = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "9" + ) + ); + final SeekableStreamPartitions endPartitions = new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "10" + ) + ); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + baseSequenceName, + startPartitions, + endPartitions, + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + while (task.getRunner().getStatus() != Status.PAUSED) { + Thread.sleep(10); + } + final Map currentOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); + + Assert.assertEquals(checkpoint1.getPartitionSequenceNumberMap(), currentOffsets); + task.getRunner().setEndOffsets(currentOffsets, false); + + while (task.getRunner().getStatus() != Status.PAUSED) { + Thread.sleep(10); + } + + final Map nextOffsets = ImmutableMap.copyOf(task.getRunner().getCurrentOffsets()); + + Assert.assertEquals(checkpoint2.getPartitionSequenceNumberMap(), nextOffsets); + + task.getRunner().setEndOffsets(nextOffsets, false); + + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + Assert.assertEquals(2, checkpointRequestsHash.size()); + Assert.assertTrue( + checkpointRequestsHash.contains( + Objects.hash( + DATA_SCHEMA.getDataSource(), + 0, + new KinesisDataSourceMetadata(startPartitions), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, currentOffsets)) + ) + ) + ); + Assert.assertTrue( + checkpointRequestsHash.contains( + Objects.hash( + DATA_SCHEMA.getDataSource(), + 0, + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, currentOffsets)), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, nextOffsets)) + ) + ) + ); + + // Check metrics + Assert.assertEquals(6, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(4, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2008/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2009/P1D", 0); + SegmentDescriptor desc3 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc4 = SD(task, "2011/P1D", 0); + SegmentDescriptor desc5 = SD(task, "2049/P1D", 0); + SegmentDescriptor desc7 = SD(task, "2013/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4, desc5, desc7), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "10" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("a"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("b"), readSegmentColumn("dim1", desc2)); + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc3)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc4)); + Assert.assertEquals(ImmutableList.of("f"), readSegmentColumn("dim1", desc5)); + Assert.assertEquals(ImmutableList.of("f"), readSegmentColumn("dim1", desc7)); + } + + + @Test(timeout = 120_000L) + public void testRunWithMinimumMessageTime() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(0, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + DateTimes.of("2010"), + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for the task to start reading + while (task.getRunner().getStatus() != Status.READING) { + Thread.sleep(10); + } + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(2, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testRunWithMaximumMessageTime() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(0, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + DateTimes.of("2010"), + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for the task to start reading + while (task.getRunner().getStatus() != Status.READING) { + Thread.sleep(10); + } + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(2, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2008/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2009/P1D", 0); + SegmentDescriptor desc3 = SD(task, "2010/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("a"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("b"), readSegmentColumn("dim1", desc2)); + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc3)); + } + + + @Test(timeout = 120_000L) + public void testRunWithTransformSpec() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(0, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + DATA_SCHEMA.withTransformSpec( + new TransformSpec( + new SelectorDimFilter("dim1", "b", null), + ImmutableList.of( + new ExpressionTransform("dim1t", "concat(dim1,dim1)", ExprMacroTable.nil()) + ) + ) + ), + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for the task to start reading + while (task.getRunner().getStatus() != Status.READING) { + Thread.sleep(10); + } + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(1, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(4, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2009/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("b"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("bb"), readSegmentColumn("dim1t", desc1)); + } + + + @Test(timeout = 120_000L) + public void testRunOnNothing() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + Assert.assertEquals(ImmutableSet.of(), publishedDescriptors()); + } + + + @Test(timeout = 60_000L) + public void testHandoffConditionTimeoutWhenHandoffOccurs() throws Exception + { + handoffConditionTimeout = 5_000; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 60_000L) + public void testHandoffConditionTimeoutWhenHandoffDoesNotOccur() throws Exception + { + doHandoff = false; + handoffConditionTimeout = 100; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testReportParseExceptions() throws Exception + { + reportParseExceptions = true; + + // these will be ignored because reportParseExceptions is true + maxParseExceptions = 1000; + maxSavedParseExceptions = 2; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "5" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.FAILED, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(1, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + Assert.assertEquals(ImmutableSet.of(), publishedDescriptors()); + Assert.assertNull(metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource())); + } + + + @Test(timeout = 120_000L) + public void testMultipleParseExceptionsSuccess() throws Exception + { + reportParseExceptions = false; + maxParseExceptions = 7; + maxSavedParseExceptions = 7; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "12" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + TaskStatus status = future.get(); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, status.getStatusCode()); + + verifyAll(); + + Assert.assertNull(status.getErrorMsg()); + + // Check metrics + Assert.assertEquals(4, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessedWithError()); + Assert.assertEquals(4, task.getRunner().getRowIngestionMeters().getUnparseable()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + SegmentDescriptor desc3 = SD(task, "2013/P1D", 0); + SegmentDescriptor desc4 = SD(task, "2049/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "12" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + IngestionStatsAndErrorsTaskReportData reportData = getTaskReportData(); + + Map expectedMetrics = ImmutableMap.of( + RowIngestionMeters.BUILD_SEGMENTS, + ImmutableMap.of( + RowIngestionMeters.PROCESSED, 4, + RowIngestionMeters.PROCESSED_WITH_ERROR, 3, + RowIngestionMeters.UNPARSEABLE, 4, + RowIngestionMeters.THROWN_AWAY, 0 + ) + ); + Assert.assertEquals(expectedMetrics, reportData.getRowStats()); + + Map unparseableEvents = ImmutableMap.of( + RowIngestionMeters.BUILD_SEGMENTS, + Arrays.asList( + "Found unparseable columns in row: [MapBasedInputRow{timestamp=2049-01-01T00:00:00.000Z, event={timestamp=2049, dim1=f, dim2=y, dimLong=10, dimFloat=20.0, met1=notanumber}, dimensions=[dim1, dim1t, dim2, dimLong, dimFloat]}], exceptions: [Unable to parse value[notanumber] for field[met1],]", + "Found unparseable columns in row: [MapBasedInputRow{timestamp=2049-01-01T00:00:00.000Z, event={timestamp=2049, dim1=f, dim2=y, dimLong=10, dimFloat=notanumber, met1=1.0}, dimensions=[dim1, dim1t, dim2, dimLong, dimFloat]}], exceptions: [could not convert value [notanumber] to float,]", + "Found unparseable columns in row: [MapBasedInputRow{timestamp=2049-01-01T00:00:00.000Z, event={timestamp=2049, dim1=f, dim2=y, dimLong=notanumber, dimFloat=20.0, met1=1.0}, dimensions=[dim1, dim1t, dim2, dimLong, dimFloat]}], exceptions: [could not convert value [notanumber] to long,]", + "Unparseable timestamp found! Event: {}", + "Unable to parse row [unparseable2]", + "Unable to parse row [unparseable]", + "Encountered row with timestamp that cannot be represented as a long: [MapBasedInputRow{timestamp=246140482-04-24T15:36:27.903Z, event={timestamp=246140482-04-24T15:36:27.903Z, dim1=x, dim2=z, dimLong=10, dimFloat=20.0, met1=1.0}, dimensions=[dim1, dim1t, dim2, dimLong, dimFloat]}]" + ) + ); + + Assert.assertEquals(unparseableEvents, reportData.getUnparseableEvents()); + } + + + @Test(timeout = 120_000L) + public void testMultipleParseExceptionsFailure() throws Exception + { + reportParseExceptions = false; + maxParseExceptions = 2; + maxSavedParseExceptions = 2; + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).once(); + + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "9" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + TaskStatus status = future.get(); + + // Wait for task to exit + Assert.assertEquals(TaskState.FAILED, status.getStatusCode()); + verifyAll(); + IndexTaskTest.checkTaskStatusErrorMsgForParseExceptionsExceeded(status); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getProcessedWithError()); + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + Assert.assertEquals(ImmutableSet.of(), publishedDescriptors()); + Assert.assertNull(metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource())); + + IngestionStatsAndErrorsTaskReportData reportData = getTaskReportData(); + + Map expectedMetrics = ImmutableMap.of( + RowIngestionMeters.BUILD_SEGMENTS, + ImmutableMap.of( + RowIngestionMeters.PROCESSED, 3, + RowIngestionMeters.PROCESSED_WITH_ERROR, 0, + RowIngestionMeters.UNPARSEABLE, 3, + RowIngestionMeters.THROWN_AWAY, 0 + ) + ); + Assert.assertEquals(expectedMetrics, reportData.getRowStats()); + + Map unparseableEvents = ImmutableMap.of( + RowIngestionMeters.BUILD_SEGMENTS, + Arrays.asList( + "Unable to parse row [unparseable2]", + "Unable to parse row [unparseable]" + ) + ); + + Assert.assertEquals(unparseableEvents, reportData.getUnparseableEvents()); + } + + + @Test(timeout = 120_000L) + public void testRunReplicas() throws Exception + { + // Insert data + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)).times(2); + + recordSupplier.close(); + expectLastCall().times(2); + + replayAll(); + + final KinesisIndexTask task1 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + final KinesisIndexTask task2 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future1 = runTask(task1); + final ListenableFuture future2 = runTask(task2); + + // Wait for tasks to exit + Assert.assertEquals(TaskState.SUCCESS, future1.get().getStatusCode()); + Assert.assertEquals(TaskState.SUCCESS, future2.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task1.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getThrownAway()); + Assert.assertEquals(3, task2.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata + SegmentDescriptor desc1 = SD(task1, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testRunConflicting() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)) + .once() + .andReturn(records.subList(3, 13)) + .once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task1 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + final KinesisIndexTask task2 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence1", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "3" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "9" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + // Run first task + final ListenableFuture future1 = runTask(task1); + Assert.assertEquals(TaskState.SUCCESS, future1.get().getStatusCode()); + + // Run second task + final ListenableFuture future2 = runTask(task2); + Assert.assertEquals(TaskState.FAILED, future2.get().getStatusCode()); + + verifyAll(); + // Check metrics + Assert.assertEquals(3, task1.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getThrownAway()); + Assert.assertEquals(3, task2.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(4, task2.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata, should all be from the first task + SegmentDescriptor desc1 = SD(task1, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + "4" + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testRunConflictingWithoutTransactions() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)) + .once() + .andReturn(records.subList(3, 13)) + .once(); + + recordSupplier.close(); + expectLastCall().times(2); + + replayAll(); + + final KinesisIndexTask task1 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + false, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + final KinesisIndexTask task2 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence1", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "3" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "9" + )), + false, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + // Run first task + final ListenableFuture future1 = runTask(task1); + Assert.assertEquals(TaskState.SUCCESS, future1.get().getStatusCode()); + + // Check published segments & metadata + SegmentDescriptor desc1 = SD(task1, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertNull(metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource())); + + // Run second task + final ListenableFuture future2 = runTask(task2); + Assert.assertEquals(TaskState.SUCCESS, future2.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task1.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getThrownAway()); + Assert.assertEquals(3, task2.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(4, task2.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata + SegmentDescriptor desc3 = SD(task2, "2011/P1D", 1); + SegmentDescriptor desc4 = SD(task2, "2013/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4), publishedDescriptors()); + Assert.assertNull(metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource())); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc3)); + Assert.assertEquals(ImmutableList.of("f"), readSegmentColumn("dim1", desc4)); + } + + + @Test(timeout = 120_000L) + public void testRunOneTaskTwoPartitions() throws Exception + { + // Insert data + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, records.size())).once(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + final KinesisIndexTask task = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence1", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2", + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4", + shardId0, + "1" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + while (countEvents(task) < 5) { + Thread.sleep(10); + } + + // Wait for tasks to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(5, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + SegmentDescriptor desc4 = SD(task, "2012/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc4), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4", + shardId0, + "1" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("g"), readSegmentColumn("dim1", desc4)); + + // Check desc2/desc3 without strong ordering because two partitions are interleaved nondeterministically + Assert.assertEquals( + ImmutableSet.of(ImmutableList.of("d", "e", "h")), + ImmutableSet.of(readSegmentColumn("dim1", desc2)) + ); + } + + + @Test(timeout = 120_000L) + public void testRunTwoTasksTwoPartitions() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)) + .once() + .andReturn(records.subList(13, 15)) + .once(); + + recordSupplier.close(); + expectLastCall().times(2); + + replayAll(); + + final KinesisIndexTask task1 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + final KinesisIndexTask task2 = createTask( + null, + new KinesisIndexTaskIOConfig( + null, + "sequence1", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future1 = runTask(task1); + Assert.assertEquals(TaskState.SUCCESS, future1.get().getStatusCode()); + + final ListenableFuture future2 = runTask(task2); + Assert.assertEquals(TaskState.SUCCESS, future2.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(3, task1.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getThrownAway()); + Assert.assertEquals(2, task2.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata + SegmentDescriptor desc1 = SD(task1, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); + SegmentDescriptor desc3 = SD(task2, "2011/P1D", 1); + SegmentDescriptor desc4 = SD(task2, "2012/P1D", 0); + + Assert.assertEquals(ImmutableSet.of(desc1, desc2, desc3, desc4), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4", + shardId0, + "1" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + // Check desc2/desc3 without strong ordering because two partitions are interleaved nondeterministically + Assert.assertEquals( + ImmutableSet.of(ImmutableList.of("d", "e"), ImmutableList.of("h")), + ImmutableSet.of(readSegmentColumn("dim1", desc2), readSegmentColumn("dim1", desc3)) + ); + Assert.assertEquals(ImmutableList.of("g"), readSegmentColumn("dim1", desc4)); + } + + + @Test(timeout = 120_000L) + public void testRestore() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 4)) + .once() + .andReturn(Collections.emptyList()) + .anyTimes(); + + replayAll(); + + final KinesisIndexTask task1 = createTask( + "task1", + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "5" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future1 = runTask(task1); + + while (countEvents(task1) != 2) { + Thread.sleep(25); + } + + Assert.assertEquals(2, countEvents(task1)); + + // Stop without publishing segment + task1.stopGracefully(); + unlockAppenderatorBasePersistDirForTask(task1); + + Assert.assertEquals(TaskState.SUCCESS, future1.get().getStatusCode()); + + verifyAll(); + reset(recordSupplier); + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(3, 6)).once(); + + recordSupplier.close(); + expectLastCall(); + + replayAll(); + + // Start a new task + final KinesisIndexTask task2 = createTask( + task1.getId(), + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "5" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + ImmutableSet.of(shardId1), + null, + null, + false + ) + ); + + final ListenableFuture future2 = runTask(task2); + + while (countEvents(task2) < 3) { + Thread.sleep(25); + } + + Assert.assertEquals(3, countEvents(task2)); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future2.get().getStatusCode()); + + verifyAll(); + + // Check metrics + Assert.assertEquals(2, task1.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task1.getRunner().getRowIngestionMeters().getThrownAway()); + Assert.assertEquals(1, task2.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(1, task2.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task2.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published segments & metadata + SegmentDescriptor desc1 = SD(task1, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task1, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "5" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 120_000L) + public void testRunWithPauseAndResume() throws Exception + { + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 5)) + .once() + .andReturn(Collections.emptyList()) + .anyTimes(); + + replayAll(); + + final KinesisIndexTask task = createTask( + "task1", + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "13" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ) + ); + + final ListenableFuture future = runTask(task); + + + while (countEvents(task) != 3) { + Thread.sleep(25); + } + + Assert.assertEquals(3, countEvents(task)); + Assert.assertEquals(Status.READING, task.getRunner().getStatus()); + + task.getRunner().pause(); + + while (task.getRunner().getStatus() != Status.PAUSED) { + Thread.sleep(10); + } + Assert.assertEquals(Status.PAUSED, task.getRunner().getStatus()); + + verifyAll(); + + Map currentOffsets = task.getRunner().getCurrentOffsets(); + + try { + future.get(10, TimeUnit.SECONDS); + Assert.fail("Task completed when it should have been paused"); + } + catch (TimeoutException e) { + // carry on.. + } + + Assert.assertEquals(currentOffsets, task.getRunner().getCurrentOffsets()); + + reset(recordSupplier); + + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + recordSupplier.close(); + expectLastCall().once(); + + replayAll(); + + task.getRunner().setEndOffsets(currentOffsets, true); + + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + verifyAll(); + Assert.assertEquals(task.getRunner().getEndOffsets(), task.getRunner().getCurrentOffsets()); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + currentOffsets.get(shardId1) + ) + )), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + + @Test(timeout = 60_000L) + public void testRunContextSequenceAheadOfStartingOffsets() throws Exception + { + // This tests the case when a replacement task is created in place of a failed test + // which has done some incremental handoffs, thus the context will contain starting + // sequence sequences from which the task should start reading and ignore the start sequences + // Insert data + recordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + + expect(recordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + + recordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(recordSupplier.poll(anyLong())).andReturn(records.subList(2, 13)) + .once(); + + recordSupplier.close(); + expectLastCall(); + + replayAll(); + + final TreeMap> sequences = new TreeMap<>(); + // Here the sequence number is 1 meaning that one incremental handoff was done by the failed task + // and this task should start reading from stream 2 for partition 0 + sequences.put(1, ImmutableMap.of(shardId1, "2")); + final Map context = new HashMap<>(); + context.put("checkpoints", objectMapper.writerWithType(new TypeReference>>() + { + }).writeValueAsString(sequences)); + + + final KinesisIndexTask task = createTask( + "task1", + new KinesisIndexTaskIOConfig( + null, + "sequence0", + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + )), + true, + null, + null, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ), + context + ); + final ListenableFuture future = runTask(task); + + // Wait for task to exit + Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); + + // Check metrics + Assert.assertEquals(3, task.getRunner().getRowIngestionMeters().getProcessed()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getUnparseable()); + Assert.assertEquals(0, task.getRunner().getRowIngestionMeters().getThrownAway()); + + // Check published metadata + SegmentDescriptor desc1 = SD(task, "2010/P1D", 0); + SegmentDescriptor desc2 = SD(task, "2011/P1D", 0); + Assert.assertEquals(ImmutableSet.of(desc1, desc2), publishedDescriptors()); + Assert.assertEquals( + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "4" + ))), + metadataStorageCoordinator.getDataSourceMetadata(DATA_SCHEMA.getDataSource()) + ); + + // Check segments in deep storage + Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", desc1)); + Assert.assertEquals(ImmutableList.of("d", "e"), readSegmentColumn("dim1", desc2)); + } + + private ListenableFuture runTask(final Task task) + { + try { + taskStorage.insert(task, TaskStatus.running(task.getId())); + } + catch (EntryExistsException e) { + // suppress + } + taskLockbox.syncFromStorage(); + final TaskToolbox toolbox = toolboxFactory.build(task); + synchronized (runningTasks) { + runningTasks.add(task); + } + return taskExec.submit( + () -> { + try { + if (task.isReady(toolbox.getTaskActionClient())) { + return task.run(toolbox); + } else { + throw new ISE("Task is not ready"); + } + } + catch (Exception e) { + log.warn(e, "Task failed"); + return TaskStatus.failure(task.getId(), Throwables.getStackTraceAsString(e)); + } + } + ); + } + + + private TaskLock getLock(final Task task, final Interval interval) + { + return Iterables.find( + taskLockbox.findLocksForTask(task), + new Predicate() + { + @Override + public boolean apply(TaskLock lock) + { + return lock.getInterval().contains(interval); + } + } + ); + } + + private KinesisIndexTask createTask( + final String taskId, + final KinesisIndexTaskIOConfig ioConfig + ) + { + return createTask(taskId, DATA_SCHEMA, ioConfig); + } + + private KinesisIndexTask createTask( + final String taskId, + final KinesisIndexTaskIOConfig ioConfig, + final Map context + ) + { + return createTask(taskId, DATA_SCHEMA, ioConfig, context); + } + + private KinesisIndexTask createTask( + final String taskId, + final DataSchema dataSchema, + final KinesisIndexTaskIOConfig ioConfig + ) + { + final KinesisIndexTaskTuningConfig tuningConfig = new KinesisIndexTaskTuningConfig( + 1000, + null, + maxRowsPerSegment, + maxTotalRows, + new Period("P1Y"), + null, + null, + null, + true, + reportParseExceptions, + handoffConditionTimeout, + resetOffsetAutomatically, + skipAvailabilityCheck, + null, + null, + null, + 5000, + null, + null, + logParseExceptions, + maxParseExceptions, + maxSavedParseExceptions, + maxRecordsPerPoll, + intermediateHandoffPeriod + ); + final Map context = null; + final KinesisIndexTask task = new TestableKinesisIndexTask( + taskId, + null, + cloneDataSchema(dataSchema), + tuningConfig, + ioConfig, + context, + null, + null, + rowIngestionMetersFactory, + null + ); + + return task; + } + + + private KinesisIndexTask createTask( + final String taskId, + final DataSchema dataSchema, + final KinesisIndexTaskIOConfig ioConfig, + final Map context + ) + { + final KinesisIndexTaskTuningConfig tuningConfig = new KinesisIndexTaskTuningConfig( + maxRowsInMemory, + null, + maxRowsPerSegment, + maxTotalRows, + new Period("P1Y"), + null, + null, + null, + true, + reportParseExceptions, + handoffConditionTimeout, + resetOffsetAutomatically, + true, + null, + null, + null, + null, + null, + null, + logParseExceptions, + maxParseExceptions, + maxSavedParseExceptions, + maxRecordsPerPoll, + intermediateHandoffPeriod + ); + context.put(SeekableStreamSupervisor.IS_INCREMENTAL_HANDOFF_SUPPORTED, true); + + final KinesisIndexTask task = new TestableKinesisIndexTask( + taskId, + null, + cloneDataSchema(dataSchema), + tuningConfig, + ioConfig, + context, + null, + null, + rowIngestionMetersFactory, + null + ); + + return task; + } + + private static DataSchema cloneDataSchema(final DataSchema dataSchema) + { + return new DataSchema( + dataSchema.getDataSource(), + dataSchema.getParserMap(), + dataSchema.getAggregators(), + dataSchema.getGranularitySpec(), + dataSchema.getTransformSpec(), + objectMapper + ); + } + + private QueryRunnerFactoryConglomerate makeTimeseriesOnlyConglomerate() + { + IntervalChunkingQueryRunnerDecorator queryRunnerDecorator = new IntervalChunkingQueryRunnerDecorator( + null, + null, + null + ) + { + @Override + public QueryRunner decorate( + QueryRunner delegate, + QueryToolChest> toolChest + ) + { + return delegate; + } + }; + return new DefaultQueryRunnerFactoryConglomerate( + ImmutableMap.of( + TimeseriesQuery.class, + new TimeseriesQueryRunnerFactory( + new TimeseriesQueryQueryToolChest(queryRunnerDecorator), + new TimeseriesQueryEngine(), + (query, future) -> { + // do nothing + } + ) + ) + ); + } + + private void makeToolboxFactory() throws IOException + { + directory = tempFolder.newFolder(); + final TestUtils testUtils = new TestUtils(); + rowIngestionMetersFactory = testUtils.getRowIngestionMetersFactory(); + final ObjectMapper objectMapper = testUtils.getTestObjectMapper(); + objectMapper.setInjectableValues(((InjectableValues.Std) objectMapper.getInjectableValues()).addValue( + AWSCredentialsConfig.class, + new AWSCredentialsConfig() + )); + for (Module module : new KinesisIndexingServiceModule().getJacksonModules()) { + objectMapper.registerModule(module); + } + final TaskConfig taskConfig = new TaskConfig( + new File(directory, "taskBaseDir").getPath(), + null, + null, + 50000, + null, + false, + null, + null + ); + final TestDerbyConnector derbyConnector = derby.getConnector(); + derbyConnector.createDataSourceTable(); + derbyConnector.createPendingSegmentsTable(); + derbyConnector.createSegmentTable(); + derbyConnector.createRulesTable(); + derbyConnector.createConfigTable(); + derbyConnector.createTaskTables(); + derbyConnector.createAuditTable(); + taskStorage = new MetadataTaskStorage( + derbyConnector, + new TaskStorageConfig(null), + new DerbyMetadataStorageActionHandlerFactory( + derbyConnector, + derby.metadataTablesConfigSupplier().get(), + objectMapper + ) + ); + metadataStorageCoordinator = new IndexerSQLMetadataStorageCoordinator( + testUtils.getTestObjectMapper(), + derby.metadataTablesConfigSupplier().get(), + derbyConnector + ); + taskLockbox = new TaskLockbox(taskStorage); + final TaskActionToolbox taskActionToolbox = new TaskActionToolbox( + taskLockbox, + taskStorage, + metadataStorageCoordinator, + emitter, + new SupervisorManager(null) + { + @Override + public boolean checkPointDataSourceMetadata( + String supervisorId, + @Nullable Integer taskGroupId, + String baseSequenceName, + @Nullable DataSourceMetadata previousDataSourceMetadata, + @Nullable DataSourceMetadata currentDataSourceMetadata + ) + { + log.info("Adding checkpoint hash to the set"); + checkpointRequestsHash.add( + Objects.hash( + supervisorId, + taskGroupId, + previousDataSourceMetadata, + currentDataSourceMetadata + ) + ); + return true; + } + }, + new Counters() + ); + final TaskActionClientFactory taskActionClientFactory = new LocalTaskActionClientFactory( + taskStorage, + taskActionToolbox, + new TaskAuditLogConfig(false) + ); + + final SegmentHandoffNotifierFactory handoffNotifierFactory = dataSource -> new SegmentHandoffNotifier() + { + @Override + public boolean registerSegmentHandoffCallback( + SegmentDescriptor descriptor, + Executor exec, + Runnable handOffRunnable + ) + { + if (doHandoff) { + // Simulate immediate handoff + exec.execute(handOffRunnable); + } + return true; + } + + @Override + public void start() + { + //Noop + } + + @Override + public void close() + { + //Noop + } + }; + final LocalDataSegmentPusherConfig dataSegmentPusherConfig = new LocalDataSegmentPusherConfig(); + dataSegmentPusherConfig.storageDirectory = getSegmentDirectory(); + final DataSegmentPusher dataSegmentPusher = new LocalDataSegmentPusher(dataSegmentPusherConfig, objectMapper); + SegmentLoaderConfig segmentLoaderConfig = new SegmentLoaderConfig() + { + @Override + public List getLocations() + { + return new ArrayList<>(); + } + }; + toolboxFactory = new TaskToolboxFactory( + taskConfig, + taskActionClientFactory, + emitter, + dataSegmentPusher, + new TestDataSegmentKiller(), + null, // DataSegmentMover + null, // DataSegmentArchiver + new TestDataSegmentAnnouncer(), + EasyMock.createNiceMock(DataSegmentServerAnnouncer.class), + handoffNotifierFactory, + this::makeTimeseriesOnlyConglomerate, + MoreExecutors.sameThreadExecutor(), // queryExecutorService + EasyMock.createMock(MonitorScheduler.class), + new SegmentLoaderFactory( + new SegmentLoaderLocalCacheManager(null, segmentLoaderConfig, testUtils.getTestObjectMapper()) + ), + testUtils.getTestObjectMapper(), + testUtils.getTestIndexIO(), + MapCache.create(1024), + new CacheConfig(), + new CachePopulatorStats(), + testUtils.getTestIndexMergerV9(), + EasyMock.createNiceMock(DruidNodeAnnouncer.class), + EasyMock.createNiceMock(DruidNode.class), + new LookupNodeService("tier"), + new DataNodeService("tier", 1, ServerType.INDEXER_EXECUTOR, 0), + new TaskReportFileWriter(reportsFile) + ); + } + + private void destroyToolboxFactory() + { + toolboxFactory = null; + taskStorage = null; + taskLockbox = null; + metadataStorageCoordinator = null; + } + + + private Set publishedDescriptors() + { + return FluentIterable.from( + metadataStorageCoordinator.getUsedSegmentsForInterval( + DATA_SCHEMA.getDataSource(), + Intervals.of("0000/3000") + ) + ).transform(DataSegment::toDescriptor).toSet(); + } + + private void unlockAppenderatorBasePersistDirForTask(KinesisIndexTask task) + throws NoSuchMethodException, InvocationTargetException, IllegalAccessException + { + Method unlockBasePersistDir = ((AppenderatorImpl) task.getAppenderator()).getClass() + .getDeclaredMethod( + "unlockBasePersistDirectory"); + unlockBasePersistDir.setAccessible(true); + unlockBasePersistDir.invoke(task.getAppenderator()); + } + + private File getSegmentDirectory() + { + return new File(directory, "segments"); + } + + + private List readSegmentColumn(final String column, final SegmentDescriptor descriptor) throws IOException + { + File indexBasePath = new File( + StringUtils.format( + "%s/%s/%s_%s/%s/%d", + getSegmentDirectory(), + DATA_SCHEMA.getDataSource(), + descriptor.getInterval().getStart(), + descriptor.getInterval().getEnd(), + descriptor.getVersion(), + descriptor.getPartitionNumber() + ) + ); + + File outputLocation = new File( + directory, + StringUtils.format( + "%s_%s_%s_%s", + descriptor.getInterval().getStart(), + descriptor.getInterval().getEnd(), + descriptor.getVersion(), + descriptor.getPartitionNumber() + ) + ); + outputLocation.mkdir(); + CompressionUtils.unzip( + Files.asByteSource(new File(indexBasePath.listFiles()[0], "index.zip")), + outputLocation, + Predicates.alwaysFalse(), + false + ); + IndexIO indexIO = new TestUtils().getTestIndexIO(); + QueryableIndex index = indexIO.loadIndex(outputLocation); + DictionaryEncodedColumn theColumn = (DictionaryEncodedColumn) index.getColumnHolder(column) + .getColumn(); + List values = new ArrayList<>(); + for (int i = 0; i < theColumn.length(); i++) { + int id = theColumn.getSingleValueRow(i); + String value = theColumn.lookupName(id); + values.add(value); + } + return values; + } + + public long countEvents(final Task task) + { + // Do a query. + TimeseriesQuery query = Druids.newTimeseriesQueryBuilder() + .dataSource(DATA_SCHEMA.getDataSource()) + .aggregators( + ImmutableList.of( + new LongSumAggregatorFactory("rows", "rows") + ) + ).granularity(Granularities.ALL) + .intervals("0000/3000") + .build(); + + List> results = + task.getQueryRunner(query).run(QueryPlus.wrap(query), ImmutableMap.of()).toList(); + + return results.isEmpty() ? 0L : DimensionHandlerUtils.nullToZero(results.get(0).getValue().getLongMetric("rows")); + } + + private static List JB( + String timestamp, + String dim1, + String dim2, + String dimLong, + String dimFloat, + String met1 + ) + { + try { + return Collections.singletonList(new ObjectMapper().writeValueAsBytes( + ImmutableMap.builder() + .put("timestamp", timestamp) + .put("dim1", dim1) + .put("dim2", dim2) + .put("dimLong", dimLong) + .put("dimFloat", dimFloat) + .put("met1", met1) + .build() + )); + } + catch (Exception e) { + throw Throwables.propagate(e); + } + } + + private SegmentDescriptor SD(final Task task, final String intervalString, final int partitionNum) + { + final Interval interval = Intervals.of(intervalString); + return new SegmentDescriptor(interval, getLock(task, interval).getVersion(), partitionNum); + } + + private IngestionStatsAndErrorsTaskReportData getTaskReportData() throws IOException + { + Map taskReports = objectMapper.readValue( + reportsFile, + new TypeReference>() + { + } + ); + return IngestionStatsAndErrorsTaskReportData.getPayloadFromTaskReports( + taskReports + ); + } + + @JsonTypeName("index_kinesis") + private static class TestableKinesisIndexTask extends KinesisIndexTask + { + @JsonCreator + public TestableKinesisIndexTask( + @JsonProperty("id") String id, + @JsonProperty("resource") TaskResource taskResource, + @JsonProperty("dataSchema") DataSchema dataSchema, + @JsonProperty("tuningConfig") KinesisIndexTaskTuningConfig tuningConfig, + @JsonProperty("ioConfig") KinesisIndexTaskIOConfig ioConfig, + @JsonProperty("context") Map context, + @JacksonInject ChatHandlerProvider chatHandlerProvider, + @JacksonInject AuthorizerMapper authorizerMapper, + @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory, + @JacksonInject AWSCredentialsConfig awsCredentialsConfig + ) + { + super( + id, + taskResource, + dataSchema, + tuningConfig, + ioConfig, + context, + chatHandlerProvider, + authorizerMapper, + rowIngestionMetersFactory, + awsCredentialsConfig + ); + } + + @Override + protected KinesisRecordSupplier newTaskRecordSupplier() + { + return recordSupplier; + } + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfigTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfigTest.java new file mode 100644 index 000000000000..4e967c414a9f --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTuningConfigTest.java @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.indexing.kinesis.supervisor.KinesisSupervisorTuningConfig; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.indexing.TuningConfig; +import org.hamcrest.CoreMatchers; +import org.joda.time.Period; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +import java.io.File; + +public class KinesisIndexTaskTuningConfigTest +{ + private final ObjectMapper mapper; + + public KinesisIndexTaskTuningConfigTest() + { + mapper = new DefaultObjectMapper(); + mapper.registerModules((Iterable) new KinesisIndexingServiceModule().getJacksonModules()); + } + + @Rule + public final ExpectedException exception = ExpectedException.none(); + + @Test + public void testSerdeWithDefaults() throws Exception + { + String jsonStr = "{\"type\": \"kinesis\"}"; + + KinesisIndexTaskTuningConfig config = (KinesisIndexTaskTuningConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + TuningConfig.class + ) + ), + TuningConfig.class + ); + + Assert.assertNotNull(config.getBasePersistDirectory()); + Assert.assertEquals(1000000, config.getMaxRowsInMemory()); + Assert.assertEquals(5_000_000, config.getMaxRowsPerSegment()); + Assert.assertEquals(new Period("PT10M"), config.getIntermediatePersistPeriod()); + Assert.assertEquals(0, config.getMaxPendingPersists()); + Assert.assertEquals(new IndexSpec(), config.getIndexSpec()); + Assert.assertTrue(config.getBuildV9Directly()); + Assert.assertFalse(config.isReportParseExceptions()); + Assert.assertEquals(0, config.getHandoffConditionTimeout()); + Assert.assertEquals(10000, config.getRecordBufferSize()); + Assert.assertEquals(5000, config.getRecordBufferOfferTimeout()); + Assert.assertEquals(5000, config.getRecordBufferFullWait()); + Assert.assertEquals(20000, config.getFetchSequenceNumberTimeout()); + Assert.assertNull(config.getFetchThreads()); + Assert.assertFalse(config.isSkipSequenceNumberAvailabilityCheck()); + Assert.assertFalse(config.isResetOffsetAutomatically()); + } + + @Test + public void testSerdeWithNonDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"basePersistDirectory\": \"/tmp/xxx\",\n" + + " \"maxRowsInMemory\": 100,\n" + + " \"maxRowsPerSegment\": 100,\n" + + " \"intermediatePersistPeriod\": \"PT1H\",\n" + + " \"maxPendingPersists\": 100,\n" + + " \"buildV9Directly\": true,\n" + + " \"reportParseExceptions\": true,\n" + + " \"handoffConditionTimeout\": 100,\n" + + " \"recordBufferSize\": 1000,\n" + + " \"recordBufferOfferTimeout\": 500,\n" + + " \"recordBufferFullWait\": 500,\n" + + " \"fetchSequenceNumberTimeout\": 6000,\n" + + " \"resetOffsetAutomatically\": false,\n" + + " \"skipSequenceNumberAvailabilityCheck\": true,\n" + + " \"fetchThreads\": 2\n" + + "}"; + + KinesisIndexTaskTuningConfig config = (KinesisIndexTaskTuningConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + TuningConfig.class + ) + ), + TuningConfig.class + ); + + Assert.assertEquals(new File("/tmp/xxx"), config.getBasePersistDirectory()); + Assert.assertEquals(100, config.getMaxRowsInMemory()); + Assert.assertEquals(100, config.getMaxRowsPerSegment()); + Assert.assertEquals(new Period("PT1H"), config.getIntermediatePersistPeriod()); + Assert.assertEquals(100, config.getMaxPendingPersists()); + Assert.assertTrue(config.getBuildV9Directly()); + Assert.assertTrue(config.isReportParseExceptions()); + Assert.assertEquals(100, config.getHandoffConditionTimeout()); + Assert.assertEquals(1000, config.getRecordBufferSize()); + Assert.assertEquals(500, config.getRecordBufferOfferTimeout()); + Assert.assertEquals(500, config.getRecordBufferFullWait()); + Assert.assertEquals(6000, config.getFetchSequenceNumberTimeout()); + Assert.assertEquals(2, (int) config.getFetchThreads()); + Assert.assertTrue(config.isSkipSequenceNumberAvailabilityCheck()); + Assert.assertFalse(config.isResetOffsetAutomatically()); + } + + @Test + public void testResetOffsetAndSkipSequenceNotBothTrue() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"basePersistDirectory\": \"/tmp/xxx\",\n" + + " \"maxRowsInMemory\": 100,\n" + + " \"maxRowsPerSegment\": 100,\n" + + " \"intermediatePersistPeriod\": \"PT1H\",\n" + + " \"maxPendingPersists\": 100,\n" + + " \"buildV9Directly\": true,\n" + + " \"reportParseExceptions\": true,\n" + + " \"handoffConditionTimeout\": 100,\n" + + " \"recordBufferSize\": 1000,\n" + + " \"recordBufferOfferTimeout\": 500,\n" + + " \"recordBufferFullWait\": 500,\n" + + " \"fetchSequenceNumberTimeout\": 6000,\n" + + " \"resetOffsetAutomatically\": true,\n" + + " \"skipSequenceNumberAvailabilityCheck\": true,\n" + + " \"fetchThreads\": 2\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(IllegalArgumentException.class)); + exception.expectMessage(CoreMatchers.containsString( + "resetOffsetAutomatically cannot be used if skipSequenceNumberAvailabilityCheck=true")); + mapper.readValue(jsonStr, TuningConfig.class); + } + + @Test + public void testConvert() + { + KinesisSupervisorTuningConfig original = new KinesisSupervisorTuningConfig( + 1, + (long) 3, + 2, + 100L, + new Period("PT3S"), + new File("/tmp/xxx"), + 4, + new IndexSpec(), + true, + true, + 5L, + true, + false, + null, + null, + null, + null, + null, + null, + 1000, + 500, + 500, + 6000, + 2, + null, + null, + null, + null, + null + ); + KinesisIndexTaskTuningConfig copy = (KinesisIndexTaskTuningConfig) original.convertToTaskTuningConfig(); + + Assert.assertEquals(1, copy.getMaxRowsInMemory()); + Assert.assertEquals(3, copy.getMaxBytesInMemory()); + Assert.assertEquals(2, copy.getMaxRowsPerSegment()); + Assert.assertEquals(100L, (long) copy.getMaxTotalRows()); + Assert.assertEquals(new Period("PT3S"), copy.getIntermediatePersistPeriod()); + Assert.assertEquals(new File("/tmp/xxx"), copy.getBasePersistDirectory()); + Assert.assertEquals(4, copy.getMaxPendingPersists()); + Assert.assertEquals(new IndexSpec(), copy.getIndexSpec()); + Assert.assertTrue(copy.getBuildV9Directly()); + Assert.assertTrue(copy.isReportParseExceptions()); + Assert.assertEquals(5L, copy.getHandoffConditionTimeout()); + Assert.assertEquals(1000, copy.getRecordBufferSize()); + Assert.assertEquals(500, copy.getRecordBufferOfferTimeout()); + Assert.assertEquals(500, copy.getRecordBufferFullWait()); + Assert.assertEquals(6000, copy.getFetchSequenceNumberTimeout()); + Assert.assertEquals(2, (int) copy.getFetchThreads()); + Assert.assertFalse(copy.isSkipSequenceNumberAvailabilityCheck()); + Assert.assertTrue(copy.isResetOffsetAutomatically()); + Assert.assertEquals(100, copy.getMaxRecordsPerPoll()); + Assert.assertEquals(new Period().withDays(Integer.MAX_VALUE), copy.getIntermediateHandoffPeriod()); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplierTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplierTest.java new file mode 100644 index 000000000000..166678cb7eb4 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRecordSupplierTest.java @@ -0,0 +1,615 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.amazonaws.services.kinesis.AmazonKinesis; +import com.amazonaws.services.kinesis.AmazonKinesisClient; +import com.amazonaws.services.kinesis.model.DescribeStreamResult; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.GetShardIteratorResult; +import com.amazonaws.services.kinesis.model.Record; +import com.amazonaws.services.kinesis.model.Shard; +import com.amazonaws.services.kinesis.model.StreamDescription; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; +import org.easymock.Capture; +import org.easymock.EasyMockSupport; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.easymock.EasyMock.anyObject; +import static org.easymock.EasyMock.anyString; +import static org.easymock.EasyMock.capture; +import static org.easymock.EasyMock.eq; +import static org.easymock.EasyMock.expect; + +public class KinesisRecordSupplierTest extends EasyMockSupport +{ + private static final String stream = "stream"; + private static long poll_timeout_millis = 2000; + private static int recordsPerFetch; + private static String shardId1 = "1"; + private static String shardId0 = "0"; + private static String shard1Iterator = "1"; + private static String shard0Iterator = "0"; + private static AmazonKinesis kinesis; + private static DescribeStreamResult describeStreamResult; + private static GetShardIteratorResult getShardIteratorResult0; + private static GetShardIteratorResult getShardIteratorResult1; + private static GetRecordsResult getRecordsResult0; + private static GetRecordsResult getRecordsResult1; + private static StreamDescription streamDescription; + private static Shard shard0; + private static Shard shard1; + private static KinesisRecordSupplier recordSupplier; + private static List shard1Records = ImmutableList.of( + new Record().withData(JB("2011", "d", "y", "10", "20.0", "1.0")).withSequenceNumber("0"), + new Record().withData(JB("2011", "e", "y", "10", "20.0", "1.0")).withSequenceNumber("1"), + new Record().withData(JB("246140482-04-24T15:36:27.903Z", "x", "z", "10", "20.0", "1.0")).withSequenceNumber("2"), + new Record().withData(ByteBuffer.wrap(StringUtils.toUtf8("unparseable"))).withSequenceNumber("3"), + new Record().withData(ByteBuffer.wrap(StringUtils.toUtf8("unparseable2"))).withSequenceNumber("4"), + new Record().withData(ByteBuffer.wrap(StringUtils.toUtf8("{}"))).withSequenceNumber("5"), + new Record().withData(JB("2013", "f", "y", "10", "20.0", "1.0")).withSequenceNumber("6"), + new Record().withData(JB("2049", "f", "y", "notanumber", "20.0", "1.0")).withSequenceNumber("7"), + new Record().withData(JB("2012", "g", "y", "10", "20.0", "1.0")).withSequenceNumber("8"), + new Record().withData(JB("2011", "h", "y", "10", "20.0", "1.0")).withSequenceNumber("9") + ); + private static List shard0Records = ImmutableList.of( + new Record().withData(JB("2008", "a", "y", "10", "20.0", "1.0")).withSequenceNumber("0"), + new Record().withData(JB("2009", "b", "y", "10", "20.0", "1.0")).withSequenceNumber("1") + ); + private static List allRecords = ImmutableList.builder() + .addAll(shard0Records.stream() + .map(x -> new OrderedPartitionableRecord<>( + stream, + shardId0, + x.getSequenceNumber(), + Collections + .singletonList( + toByteArray( + x.getData())) + )) + .collect( + Collectors + .toList())) + .addAll(shard1Records.stream() + .map(x -> new OrderedPartitionableRecord<>( + stream, + shardId1, + x.getSequenceNumber(), + Collections + .singletonList( + toByteArray( + x.getData())) + )) + .collect( + Collectors + .toList())) + .build(); + + private static ByteBuffer JB(String timestamp, String dim1, String dim2, String dimLong, String dimFloat, String met1) + { + try { + return ByteBuffer.wrap(new ObjectMapper().writeValueAsBytes( + ImmutableMap.builder() + .put("timestamp", timestamp) + .put("dim1", dim1) + .put("dim2", dim2) + .put("dimLong", dimLong) + .put("dimFloat", dimFloat) + .put("met1", met1) + .build() + )); + } + catch (Exception e) { + throw Throwables.propagate(e); + } + } + + @Before + public void setupTest() + { + kinesis = createMock(AmazonKinesisClient.class); + describeStreamResult = createMock(DescribeStreamResult.class); + getShardIteratorResult0 = createMock(GetShardIteratorResult.class); + getShardIteratorResult1 = createMock(GetShardIteratorResult.class); + getRecordsResult0 = createMock(GetRecordsResult.class); + getRecordsResult1 = createMock(GetRecordsResult.class); + streamDescription = createMock(StreamDescription.class); + shard0 = createMock(Shard.class); + shard1 = createMock(Shard.class); + recordsPerFetch = 1; + } + + @After + public void tearDownTest() + { + recordSupplier.close(); + recordSupplier = null; + } + + @Test + public void testSupplierSetup() + { + Capture captured = Capture.newInstance(); + expect(kinesis.describeStream(capture(captured))).andReturn(describeStreamResult).once(); + expect(describeStreamResult.getStreamDescription()).andReturn(streamDescription).once(); + expect(streamDescription.getShards()).andReturn(ImmutableList.of(shard0, shard1)).once(); + expect(shard0.getShardId()).andReturn(shardId0).once(); + expect(shard1.getShardId()).andReturn(shardId1).once(); + + replayAll(); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(stream, shardId0), + StreamPartition.of(stream, shardId1) + ); + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 5 + ); + + Assert.assertTrue(recordSupplier.getAssignment().isEmpty()); + + recordSupplier.assign(partitions); + + Assert.assertEquals(partitions, recordSupplier.getAssignment()); + Assert.assertEquals(ImmutableSet.of(shardId1, shardId0), recordSupplier.getPartitionIds(stream)); + Assert.assertEquals(Collections.emptyList(), recordSupplier.poll(100)); + + verifyAll(); + Assert.assertEquals(stream, captured.getValue()); + } + + private static GetRecordsRequest generateGetRecordsReq(String shardIterator, int limit) + { + return new GetRecordsRequest().withShardIterator(shardIterator).withLimit(limit); + } + + // filter out EOS markers + private static List> cleanRecords(List> records) + { + return records.stream() + .filter(x -> !x.getSequenceNumber() + .equals(KinesisSequenceNumber.END_OF_SHARD_MARKER)) + .collect(Collectors.toList()); + } + + @Test + public void testPoll() throws InterruptedException + { + recordsPerFetch = 100; + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId0), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult0).anyTimes(); + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId1), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult1).anyTimes(); + + expect(getShardIteratorResult0.getShardIterator()).andReturn(shard0Iterator).anyTimes(); + expect(getShardIteratorResult1.getShardIterator()).andReturn(shard1Iterator).anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard0Iterator, recordsPerFetch))).andReturn(getRecordsResult0) + .anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard1Iterator, recordsPerFetch))).andReturn(getRecordsResult1) + .anyTimes(); + expect(getRecordsResult0.getRecords()).andReturn(shard0Records).once(); + expect(getRecordsResult1.getRecords()).andReturn(shard1Records).once(); + expect(getRecordsResult0.getNextShardIterator()).andReturn(null).anyTimes(); + expect(getRecordsResult1.getNextShardIterator()).andReturn(null).anyTimes(); + + replayAll(); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(stream, shardId0), + StreamPartition.of(stream, shardId1) + ); + + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 100 + ); + + recordSupplier.assign(partitions); + recordSupplier.seekToEarliest(partitions); + recordSupplier.start(); + + while (recordSupplier.bufferSize() < 12) { + Thread.sleep(100); + } + + List> polledRecords = cleanRecords(recordSupplier.poll( + poll_timeout_millis)); + + verifyAll(); + + Assert.assertEquals(partitions, recordSupplier.getAssignment()); + Assert.assertTrue(polledRecords.containsAll(allRecords)); + } + + @Test + public void testSeek() + throws InterruptedException + { + recordsPerFetch = 100; + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId0), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult0).anyTimes(); + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId1), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult1).anyTimes(); + + expect(getShardIteratorResult0.getShardIterator()).andReturn(shard0Iterator).anyTimes(); + expect(getShardIteratorResult1.getShardIterator()).andReturn(shard1Iterator).anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard0Iterator, recordsPerFetch))).andReturn(getRecordsResult0) + .anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard1Iterator, recordsPerFetch))).andReturn(getRecordsResult1) + .anyTimes(); + expect(getRecordsResult0.getRecords()).andReturn(shard0Records.subList(1, shard0Records.size())).once(); + expect(getRecordsResult1.getRecords()).andReturn(shard1Records.subList(2, shard1Records.size())).once(); + expect(getRecordsResult0.getNextShardIterator()).andReturn(null).anyTimes(); + expect(getRecordsResult1.getNextShardIterator()).andReturn(null).anyTimes(); + + replayAll(); + + StreamPartition shard0Partition = StreamPartition.of(stream, shardId0); + StreamPartition shard1Partition = StreamPartition.of(stream, shardId1); + Set> partitions = ImmutableSet.of( + shard0Partition, + shard1Partition + ); + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 100 + ); + + recordSupplier.assign(partitions); + recordSupplier.seek(shard1Partition, shard1Records.get(2).getSequenceNumber()); + recordSupplier.seek(shard0Partition, shard0Records.get(1).getSequenceNumber()); + recordSupplier.start(); + + for (int i = 0; i < 10 && recordSupplier.bufferSize() < 9; i++) { + Thread.sleep(100); + } + + List> polledRecords = cleanRecords(recordSupplier.poll( + poll_timeout_millis)); + + verifyAll(); + Assert.assertEquals(9, polledRecords.size()); + Assert.assertTrue(polledRecords.containsAll(allRecords.subList(4, 12))); + Assert.assertTrue(polledRecords.containsAll(allRecords.subList(1, 2))); + + } + + + @Test + public void testSeekToLatest() + throws InterruptedException + { + recordsPerFetch = 100; + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId0), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult0).anyTimes(); + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId1), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult1).anyTimes(); + + expect(getShardIteratorResult0.getShardIterator()).andReturn(null).once(); + expect(getShardIteratorResult1.getShardIterator()).andReturn(null).once(); + + replayAll(); + + StreamPartition shard0 = StreamPartition.of(stream, shardId0); + StreamPartition shard1 = StreamPartition.of(stream, shardId1); + Set> partitions = ImmutableSet.of( + shard0, + shard1 + ); + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 100 + ); + + recordSupplier.assign(partitions); + recordSupplier.seekToLatest(partitions); + recordSupplier.start(); + + for (int i = 0; i < 10 && recordSupplier.bufferSize() < 2; i++) { + Thread.sleep(100); + } + Assert.assertEquals(Collections.emptyList(), cleanRecords(recordSupplier.poll(poll_timeout_millis))); + + verifyAll(); + } + + @Test(expected = ISE.class) + public void testSeekUnassigned() throws InterruptedException + { + StreamPartition shard0 = StreamPartition.of(stream, shardId0); + StreamPartition shard1 = StreamPartition.of(stream, shardId1); + Set> partitions = ImmutableSet.of( + shard1 + ); + + recordSupplier = new KinesisRecordSupplier( + kinesis, + 1, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 5 + ); + + recordSupplier.assign(partitions); + recordSupplier.seekToEarliest(Collections.singleton(shard0)); + } + + + @Test + public void testPollAfterSeek() + throws InterruptedException + { + // tests that after doing a seek, the now invalid records in buffer is cleaned up properly + recordsPerFetch = 100; + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId1), + anyString(), + eq("5") + )).andReturn( + getShardIteratorResult1).once(); + + expect(kinesis.getShardIterator(anyObject(), eq(shardId1), anyString(), eq("7"))).andReturn(getShardIteratorResult0) + .once(); + + expect(getShardIteratorResult1.getShardIterator()).andReturn(shard1Iterator).once(); + expect(getShardIteratorResult0.getShardIterator()).andReturn(shard0Iterator).once(); + expect(kinesis.getRecords(generateGetRecordsReq(shard1Iterator, recordsPerFetch))).andReturn(getRecordsResult1) + .once(); + expect(kinesis.getRecords(generateGetRecordsReq(shard0Iterator, recordsPerFetch))).andReturn(getRecordsResult0) + .once(); + expect(getRecordsResult1.getRecords()).andReturn(shard1Records.subList(5, shard1Records.size())).once(); + expect(getRecordsResult0.getRecords()).andReturn(shard1Records.subList(7, shard1Records.size())).once(); + expect(getRecordsResult1.getNextShardIterator()).andReturn(null).anyTimes(); + expect(getRecordsResult0.getNextShardIterator()).andReturn(null).anyTimes(); + + replayAll(); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(stream, shardId1) + ); + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + false, + 100, + 5000, + 5000, + 60000, + 1 + ); + + recordSupplier.assign(partitions); + recordSupplier.seek(StreamPartition.of(stream, shardId1), "5"); + recordSupplier.start(); + + for (int i = 0; i < 10 && recordSupplier.bufferSize() < 6; i++) { + Thread.sleep(100); + } + + OrderedPartitionableRecord firstRecord = recordSupplier.poll(poll_timeout_millis).get(0); + + Assert.assertEquals( + allRecords.get(7), + firstRecord + ); + + recordSupplier.seek(StreamPartition.of(stream, shardId1), "7"); + recordSupplier.start(); + + while (recordSupplier.bufferSize() < 4) { + Thread.sleep(100); + } + + OrderedPartitionableRecord record2 = recordSupplier.poll(poll_timeout_millis).get(0); + + Assert.assertEquals(allRecords.get(9), record2); + verifyAll(); + } + + + @Test + public void testPollDeaggregate() throws InterruptedException + { + recordsPerFetch = 100; + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId0), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult0).anyTimes(); + + expect(kinesis.getShardIterator( + anyObject(), + eq(shardId1), + anyString(), + anyString() + )).andReturn( + getShardIteratorResult1).anyTimes(); + + expect(getShardIteratorResult0.getShardIterator()).andReturn(shard0Iterator).anyTimes(); + expect(getShardIteratorResult1.getShardIterator()).andReturn(shard1Iterator).anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard0Iterator, recordsPerFetch))).andReturn(getRecordsResult0) + .anyTimes(); + expect(kinesis.getRecords(generateGetRecordsReq(shard1Iterator, recordsPerFetch))).andReturn(getRecordsResult1) + .anyTimes(); + expect(getRecordsResult0.getRecords()).andReturn(shard0Records).once(); + expect(getRecordsResult1.getRecords()).andReturn(shard1Records).once(); + expect(getRecordsResult0.getNextShardIterator()).andReturn(null).anyTimes(); + expect(getRecordsResult1.getNextShardIterator()).andReturn(null).anyTimes(); + + replayAll(); + + Set> partitions = ImmutableSet.of( + StreamPartition.of(stream, shardId0), + StreamPartition.of(stream, shardId1) + ); + + + recordSupplier = new KinesisRecordSupplier( + kinesis, + recordsPerFetch, + 0, + 2, + true, + 100, + 5000, + 5000, + 60000, + 100 + ); + + recordSupplier.assign(partitions); + recordSupplier.seekToEarliest(partitions); + recordSupplier.start(); + + for (int i = 0; i < 10 && recordSupplier.bufferSize() < 12; i++) { + Thread.sleep(100); + } + + List> polledRecords = cleanRecords(recordSupplier.poll( + poll_timeout_millis)); + + verifyAll(); + + Assert.assertEquals(partitions, recordSupplier.getAssignment()); + Assert.assertTrue(polledRecords.containsAll(allRecords)); + } + + /** + * Returns an array with the content between the position and limit of "buffer". This may be the buffer's backing + * array itself. Does not modify position or limit of the buffer. + */ + private static byte[] toByteArray(final ByteBuffer buffer) + { + if (buffer.hasArray() + && buffer.arrayOffset() == 0 + && buffer.position() == 0 + && buffer.array().length == buffer.limit()) { + return buffer.array(); + } else { + final byte[] retVal = new byte[buffer.remaining()]; + buffer.duplicate().get(retVal); + return retVal; + } + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRegionTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRegionTest.java new file mode 100644 index 000000000000..b20cef3fc333 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisRegionTest.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.io.IOException; + +public class KinesisRegionTest +{ + private ObjectMapper mapper; + + @Before + public void setupTest() + { + mapper = new DefaultObjectMapper(); + } + + @Test + public void testSerde() throws IOException + { + KinesisRegion kinesisRegionUs1 = KinesisRegion.US_EAST_1; + KinesisRegion kinesisRegionAp1 = KinesisRegion.AP_NORTHEAST_1; + + Assert.assertEquals("\"us-east-1\"", mapper.writeValueAsString(kinesisRegionUs1)); + Assert.assertEquals("\"ap-northeast-1\"", mapper.writeValueAsString(kinesisRegionAp1)); + + KinesisRegion kinesisRegion = mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + "\"us-east-1\"", + KinesisRegion.class + ) + ), + KinesisRegion.class + ); + + Assert.assertEquals(kinesisRegion, KinesisRegion.US_EAST_1); + } + + @Test + public void testGetEndpoint() + { + Assert.assertEquals("kinesis.cn-north-1.amazonaws.com.cn", KinesisRegion.CN_NORTH_1.getEndpoint()); + Assert.assertEquals("kinesis.us-east-1.amazonaws.com", KinesisRegion.US_EAST_1.getEndpoint()); + } + + @Test(expected = JsonMappingException.class) + public void testBadSerde() throws IOException + { + mapper.readValue( + "\"us-east-10\"", + KinesisRegion.class + ); + } + +} + diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfigTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfigTest.java new file mode 100644 index 000000000000..037c92bb0511 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorIOConfigTest.java @@ -0,0 +1,139 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.databind.JsonMappingException; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.indexing.kinesis.KinesisIndexingServiceModule; +import org.apache.druid.indexing.kinesis.KinesisRegion; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.hamcrest.CoreMatchers; +import org.joda.time.Duration; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +public class KinesisSupervisorIOConfigTest +{ + private final ObjectMapper mapper; + + public KinesisSupervisorIOConfigTest() + { + mapper = new DefaultObjectMapper(); + mapper.registerModules((Iterable) new KinesisIndexingServiceModule().getJacksonModules()); + } + + @Rule + public final ExpectedException exception = ExpectedException.none(); + + @Test + public void testSerdeWithDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"stream\": \"my-stream\"\n" + + "}"; + + KinesisSupervisorIOConfig config = mapper.readValue( + jsonStr, + KinesisSupervisorIOConfig.class + ); + + Assert.assertEquals("my-stream", config.getStream()); + Assert.assertEquals(KinesisRegion.US_EAST_1.getEndpoint(), config.getEndpoint()); + Assert.assertEquals(1, (int) config.getReplicas()); + Assert.assertEquals(1, (int) config.getTaskCount()); + Assert.assertEquals(Duration.standardMinutes(60), config.getTaskDuration()); + Assert.assertEquals(Duration.standardSeconds(5), config.getStartDelay()); + Assert.assertEquals(Duration.standardSeconds(30), config.getPeriod()); + Assert.assertFalse(config.isUseEarliestSequenceNumber()); + Assert.assertEquals(Duration.standardMinutes(30), config.getCompletionTimeout()); + Assert.assertFalse("lateMessageRejectionPeriod", config.getLateMessageRejectionPeriod().isPresent()); + Assert.assertFalse("earlyMessageRejectionPeriod", config.getEarlyMessageRejectionPeriod().isPresent()); + Assert.assertEquals((Integer) 4000, config.getRecordsPerFetch()); + Assert.assertEquals((Integer) 0, config.getFetchDelayMillis()); + Assert.assertNull(config.getAwsAssumedRoleArn()); + Assert.assertNull(config.getAwsExternalId()); + Assert.assertFalse(config.isDeaggregate()); + } + + @Test + public void testSerdeWithNonDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"stream\": \"my-stream\",\n" + + " \"region\": \"us-east-2\",\n" + + " \"endpoint\": \"kinesis.us-east-2.amazonaws.com\",\n" + + " \"replicas\": 3,\n" + + " \"taskCount\": 9,\n" + + " \"taskDuration\": \"PT30M\",\n" + + " \"startDelay\": \"PT1M\",\n" + + " \"period\": \"PT10S\",\n" + + " \"useEarliestSequenceNumber\": true,\n" + + " \"completionTimeout\": \"PT45M\",\n" + + " \"lateMessageRejectionPeriod\": \"PT1H\",\n" + + " \"earlyMessageRejectionPeriod\": \"PT1H\",\n" + + " \"recordsPerFetch\": 4000,\n" + + " \"fetchDelayMillis\": 1000,\n" + + " \"awsAssumedRoleArn\": \"role\",\n" + + " \"awsExternalId\": \"awsexternalid\",\n" + + " \"deaggregate\": true\n" + + "}"; + + KinesisSupervisorIOConfig config = mapper.readValue( + jsonStr, + KinesisSupervisorIOConfig.class + ); + + Assert.assertEquals("my-stream", config.getStream()); + Assert.assertEquals(config.getEndpoint(), "kinesis.us-east-2.amazonaws.com"); + Assert.assertEquals(3, (int) config.getReplicas()); + Assert.assertEquals(9, (int) config.getTaskCount()); + Assert.assertEquals(Duration.standardMinutes(30), config.getTaskDuration()); + Assert.assertEquals(Duration.standardMinutes(1), config.getStartDelay()); + Assert.assertEquals(Duration.standardSeconds(10), config.getPeriod()); + Assert.assertTrue(config.isUseEarliestSequenceNumber()); + Assert.assertEquals(Duration.standardMinutes(45), config.getCompletionTimeout()); + Assert.assertEquals(Duration.standardHours(1), config.getLateMessageRejectionPeriod().get()); + Assert.assertEquals(Duration.standardHours(1), config.getEarlyMessageRejectionPeriod().get()); + Assert.assertEquals((Integer) 4000, config.getRecordsPerFetch()); + Assert.assertEquals((Integer) 1000, config.getFetchDelayMillis()); + Assert.assertEquals("role", config.getAwsAssumedRoleArn()); + Assert.assertEquals("awsexternalid", config.getAwsExternalId()); + Assert.assertTrue(config.isDeaggregate()); + } + + @Test + public void testTopicRequired() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\"\n" + + "}"; + + exception.expect(JsonMappingException.class); + exception.expectCause(CoreMatchers.isA(NullPointerException.class)); + exception.expectMessage(CoreMatchers.containsString("stream")); + mapper.readValue(jsonStr, KinesisSupervisorIOConfig.class); + } + +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java new file mode 100644 index 000000000000..750b9247d792 --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java @@ -0,0 +1,3785 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Optional; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.data.input.impl.DimensionSchema; +import org.apache.druid.data.input.impl.DimensionsSpec; +import org.apache.druid.data.input.impl.JSONParseSpec; +import org.apache.druid.data.input.impl.StringDimensionSchema; +import org.apache.druid.data.input.impl.StringInputRowParser; +import org.apache.druid.data.input.impl.TimestampSpec; +import org.apache.druid.indexer.TaskLocation; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.common.TestUtils; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.RealtimeIndexTask; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.kinesis.KinesisDataSourceMetadata; +import org.apache.druid.indexing.kinesis.KinesisIndexTask; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskClient; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskClientFactory; +import org.apache.druid.indexing.kinesis.KinesisIndexTaskIOConfig; +import org.apache.druid.indexing.kinesis.KinesisRecordSupplier; +import org.apache.druid.indexing.overlord.DataSourceMetadata; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskQueue; +import org.apache.druid.indexing.overlord.TaskRunner; +import org.apache.druid.indexing.overlord.TaskRunnerListener; +import org.apache.druid.indexing.overlord.TaskRunnerWorkItem; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.supervisor.SupervisorReport; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.indexing.seekablestream.supervisor.TaskReportData; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.granularity.Granularities; +import org.apache.druid.java.util.common.parsers.JSONPathSpec; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.indexing.RealtimeIOConfig; +import org.apache.druid.segment.indexing.granularity.UniformGranularitySpec; +import org.apache.druid.segment.realtime.FireDepartment; +import org.apache.druid.server.metrics.DruidMonitorSchedulerConfig; +import org.apache.druid.server.metrics.ExceptionCapturingServiceEmitter; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.easymock.Capture; +import org.easymock.CaptureType; +import org.easymock.EasyMock; +import org.easymock.EasyMockSupport; +import org.easymock.IAnswer; +import org.joda.time.DateTime; +import org.joda.time.Duration; +import org.joda.time.Period; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeoutException; + +import static org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.Status; +import static org.easymock.EasyMock.anyObject; +import static org.easymock.EasyMock.anyString; +import static org.easymock.EasyMock.expect; +import static org.easymock.EasyMock.expectLastCall; + +public class KinesisSupervisorTest extends EasyMockSupport +{ + + + private static final ObjectMapper objectMapper = TestHelper.makeJsonMapper(); + private static final String DATASOURCE = "testDS"; + private static final int TEST_CHAT_THREADS = 3; + private static final long TEST_CHAT_RETRIES = 9L; + private static final Period TEST_HTTP_TIMEOUT = new Period("PT10S"); + private static final Period TEST_SHUTDOWN_TIMEOUT = new Period("PT80S"); + private static final String stream = "stream"; + private static String shardId1 = "1"; + private static String shardId0 = "0"; + private static StreamPartition shard1Partition = StreamPartition.of(stream, shardId1); + private static StreamPartition shard0Partition = StreamPartition.of(stream, shardId0); + + private static DataSchema dataSchema; + private KinesisRecordSupplier supervisorRecordSupplier; + + private final int numThreads; + private KinesisSupervisor supervisor; + private KinesisSupervisorTuningConfig tuningConfig; + private TaskStorage taskStorage; + private TaskMaster taskMaster; + private TaskRunner taskRunner; + private IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator; + private KinesisIndexTaskClient taskClient; + private TaskQueue taskQueue; + private RowIngestionMetersFactory rowIngestionMetersFactory; + private ExceptionCapturingServiceEmitter serviceEmitter; + + public KinesisSupervisorTest() + { + this.numThreads = 1; + } + + @BeforeClass + public static void setupClass() + { + dataSchema = getDataSchema(DATASOURCE); + } + + @Before + public void setupTest() + { + taskStorage = createMock(TaskStorage.class); + taskMaster = createMock(TaskMaster.class); + taskRunner = createMock(TaskRunner.class); + indexerMetadataStorageCoordinator = createMock(IndexerMetadataStorageCoordinator.class); + taskClient = createMock(KinesisIndexTaskClient.class); + taskQueue = createMock(TaskQueue.class); + supervisorRecordSupplier = createMock(KinesisRecordSupplier.class); + + tuningConfig = new KinesisSupervisorTuningConfig( + 1000, + null, + 50000, + null, + new Period("P1Y"), + new File("/test"), + null, + null, + true, + false, + null, + null, + null, + null, + numThreads, + TEST_CHAT_THREADS, + TEST_CHAT_RETRIES, + TEST_HTTP_TIMEOUT, + TEST_SHUTDOWN_TIMEOUT, + null, + null, + null, + 5000, + null, + null, + null, + null, + null, + null + ); + rowIngestionMetersFactory = new TestUtils().getRowIngestionMetersFactory(); + serviceEmitter = new ExceptionCapturingServiceEmitter(); + EmittingLogger.registerEmitter(serviceEmitter); + } + + @After + public void tearDownTest() + { + supervisor = null; + } + + + @Test + public void testNoInitialState() throws Exception + { + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(); + expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true); + taskRunner.registerListener(anyObject(TaskRunnerListener.class), anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task = captured.getValue(); + Assert.assertEquals(dataSchema, task.getDataSchema()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), task.getTuningConfig()); + + KinesisIndexTaskIOConfig taskConfig = task.getIOConfig(); + Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); + Assert.assertTrue("isUseTransaction", taskConfig.isUseTransaction()); + Assert.assertFalse("minimumMessageTime", taskConfig.getMinimumMessageTime().isPresent()); + Assert.assertFalse("maximumMessageTime", taskConfig.getMaximumMessageTime().isPresent()); + + Assert.assertEquals(stream, taskConfig.getStartPartitions().getStream()); + Assert.assertEquals( + "0", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + "0", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + + Assert.assertEquals(stream, taskConfig.getEndPartitions().getStream()); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + taskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + + + @Test + public void testMultiTask() throws Exception + { + supervisor = getSupervisor(1, 2, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task1 = captured.getValues().get(0); + Assert.assertEquals(1, task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(1, task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals( + "0", + task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + + KinesisIndexTask task2 = captured.getValues().get(1); + Assert.assertEquals(1, task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(1, task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals( + "0", + task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + + @Test + public void testReplicas() throws Exception + { + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task1 = captured.getValues().get(0); + Assert.assertEquals(2, task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(2, task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals( + "0", + task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + Assert.assertEquals( + "0", + task1.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task1.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + + KinesisIndexTask task2 = captured.getValues().get(1); + Assert.assertEquals(2, task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals(2, task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().size()); + Assert.assertEquals( + "0", + task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + Assert.assertEquals( + "0", + task2.getIOConfig().getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + task2.getIOConfig().getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + + } + + @Test + public void testLateMessageRejectionPeriod() throws Exception + { + supervisor = getSupervisor(2, 1, true, "PT1H", new Period("PT1H"), null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task1 = captured.getValues().get(0); + KinesisIndexTask task2 = captured.getValues().get(1); + + Assert.assertTrue( + "minimumMessageTime", + task1.getIOConfig().getMinimumMessageTime().get().plusMinutes(59).isBeforeNow() + ); + Assert.assertTrue( + "minimumMessageTime", + task1.getIOConfig().getMinimumMessageTime().get().plusMinutes(61).isAfterNow() + ); + Assert.assertEquals( + task1.getIOConfig().getMinimumMessageTime().get(), + task2.getIOConfig().getMinimumMessageTime().get() + ); + } + + @Test + public void testEarlyMessageRejectionPeriod() throws Exception + { + supervisor = getSupervisor(2, 1, true, "PT1H", null, new Period("PT1H")); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task1 = captured.getValues().get(0); + KinesisIndexTask task2 = captured.getValues().get(1); + + Assert.assertTrue( + "maximumMessageTime", + task1.getIOConfig().getMaximumMessageTime().get().minusMinutes(59 + 60).isAfterNow() + ); + Assert.assertTrue( + "maximumMessageTime", + task1.getIOConfig().getMaximumMessageTime().get().minusMinutes(61 + 60).isBeforeNow() + ); + Assert.assertEquals( + task1.getIOConfig().getMaximumMessageTime().get(), + task2.getIOConfig().getMaximumMessageTime().get() + ); + } + + + @Test + /** + * Test generating the starting sequences from the partition data stored in druid_dataSource which contains the + * sequences of the last built segments. + */ + public void testDatasourceMetadata() throws Exception + { + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(shard1Partition)).andReturn("2").anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )) + ) + ).anyTimes(); + + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisIndexTask task = captured.getValue(); + KinesisIndexTaskIOConfig taskConfig = task.getIOConfig(); + Assert.assertEquals("sequenceName-0", taskConfig.getBaseSequenceName()); + Assert.assertEquals( + "2", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + "1", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + + @Test(expected = ISE.class) + public void testBadMetadataOffsets() throws Exception + { + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "101", + shardId0, + "-1" + )) + ) + ).anyTimes(); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + } + + @Test + public void testKillIncompatibleTasks() throws Exception + { + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + // unexpected # of partitions (kill) + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 1, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + null, + null + ); + + // correct number of partitions and ranges (don't kill) + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0", + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1", + shardId1, + "12" + )), + null, + null + ); + + // unexpected range on partition 2 (kill) + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 1, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0", + shardId1, + "1" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1", + shardId1, + "11" + )), + null, + null + ); + + // different datasource (don't kill) + Task id4 = createKinesisIndexTask( + "id4", + "other-datasource", + 2, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0", + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1", + shardId1, + "12" + )), + null, + null + ); + + // non KinesisIndexTask (don't kill) + Task id5 = new RealtimeIndexTask( + "id5", + null, + new FireDepartment( + dataSchema, + new RealtimeIOConfig(null, null, null), + null + ), + null + ); + + List existingTasks = ImmutableList.of(id1, id2, id3, id4, id5); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.stopAsync("id1", false)).andReturn(Futures.immediateFuture(true)); + EasyMock.expect(taskClient.stopAsync("id3", false)).andReturn(Futures.immediateFuture(false)); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + taskQueue.shutdown("id3", "Task [%s] failed to stop in a timely manner, killing task", "id3"); + + EasyMock.expect(taskQueue.add(EasyMock.anyObject(Task.class))).andReturn(true); + + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId0, + "0", + shardId1, + "0" + )); + + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(2); + + replayAll(); + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + } + + @Test + public void testKillBadPartitionAssignment() throws Exception + { + supervisor = getSupervisor(1, 2, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId1, + "12" + )), + null, + null + ); + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 1, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + null, + null + ); + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0", + shardId1, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1", + shardId1, + "12" + )), + null, + null + ); + Task id4 = createKinesisIndexTask( + "id4", + DATASOURCE, + 0, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + null, + null + ); + Task id5 = createKinesisIndexTask( + "id5", + DATASOURCE, + 0, + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "0" + )), + new SeekableStreamPartitions<>(stream, ImmutableMap.of( + shardId0, + "1" + )), + null, + null + ); + + List existingTasks = ImmutableList.of(id1, id2, id3, id4, id5); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id4")).andReturn(Optional.of(TaskStatus.running("id4"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id5")).andReturn(Optional.of(TaskStatus.running("id5"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id4")).andReturn(Optional.of(id4)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id5")).andReturn(Optional.of(id5)).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.stopAsync("id3", false)).andReturn(Futures.immediateFuture(true)); + EasyMock.expect(taskClient.stopAsync("id4", false)).andReturn(Futures.immediateFuture(false)); + EasyMock.expect(taskClient.stopAsync("id5", false)).andReturn(Futures.immediateFuture((Boolean) null)); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of(shardId1, "0")); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of(shardId0, "0")); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(1); + + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + taskQueue.shutdown("id4", "Task [%s] failed to stop in a timely manner, killing task", "id4"); + taskQueue.shutdown("id5", "Task [%s] failed to stop in a timely manner, killing task", "id5"); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + } + + @Test + public void testRequeueTaskWhenFailed() throws Exception + { + supervisor = getSupervisor(2, 2, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put( + 0, + ImmutableMap.of( + shardId1, + "0" + ) + ); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of( + shardId0, + "0" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .anyTimes(); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .anyTimes(); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + + // test that running the main loop again checks the status of the tasks that were created and does nothing if they + // are all still running + EasyMock.reset(taskStorage); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.replay(taskStorage); + + supervisor.runInternal(); + verifyAll(); + + // test that a task failing causes a new task to be re-queued with the same parameters + Capture aNewTaskCapture = Capture.newInstance(); + List imStillAlive = tasks.subList(0, 3); + KinesisIndexTask iHaveFailed = (KinesisIndexTask) tasks.get(3); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(imStillAlive).anyTimes(); + for (Task task : imStillAlive) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.expect(taskStorage.getStatus(iHaveFailed.getId())) + .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); + EasyMock.expect(taskStorage.getTask(iHaveFailed.getId())).andReturn(Optional.of((Task) iHaveFailed)).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(aNewTaskCapture))).andReturn(true); + EasyMock.replay(taskStorage); + EasyMock.replay(taskQueue); + + supervisor.runInternal(); + verifyAll(); + + Assert.assertNotEquals(iHaveFailed.getId(), aNewTaskCapture.getValue().getId()); + Assert.assertEquals( + iHaveFailed.getIOConfig().getBaseSequenceName(), + ((KinesisIndexTask) aNewTaskCapture.getValue()).getIOConfig().getBaseSequenceName() + ); + } + + @Test + public void testRequeueAdoptedTaskWhenFailed() throws Exception + { + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + DateTime now = DateTimes.nowUtc(); + DateTime maxi = now.plusMinutes(60); + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + now, + maxi + ); + + List existingTasks = ImmutableList.of(id1); + + Capture captured = Capture.newInstance(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(existingTasks).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(now)).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(2); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + // check that replica tasks are created with the same minimumMessageTime as tasks inherited from another supervisor + Assert.assertEquals(now, ((KinesisIndexTask) captured.getValue()).getIOConfig().getMinimumMessageTime().get()); + + // test that a task failing causes a new task to be re-queued with the same parameters + String runningTaskId = captured.getValue().getId(); + Capture aNewTaskCapture = Capture.newInstance(); + KinesisIndexTask iHaveFailed = (KinesisIndexTask) existingTasks.get(0); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); + EasyMock.reset(taskClient); + + // for the newly created replica task + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(captured.getValue())).anyTimes(); + EasyMock.expect(taskStorage.getStatus(iHaveFailed.getId())) + .andReturn(Optional.of(TaskStatus.failure(iHaveFailed.getId()))); + EasyMock.expect(taskStorage.getStatus(runningTaskId)) + .andReturn(Optional.of(TaskStatus.running(runningTaskId))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(iHaveFailed.getId())).andReturn(Optional.of((Task) iHaveFailed)).anyTimes(); + EasyMock.expect(taskStorage.getTask(runningTaskId)).andReturn(Optional.of(captured.getValue())).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(runningTaskId)) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync(runningTaskId)).andReturn(Futures.immediateFuture(now)).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(aNewTaskCapture))).andReturn(true); + EasyMock.replay(taskStorage); + EasyMock.replay(taskQueue); + EasyMock.replay(taskClient); + + supervisor.runInternal(); + verifyAll(); + + Assert.assertNotEquals(iHaveFailed.getId(), aNewTaskCapture.getValue().getId()); + Assert.assertEquals( + iHaveFailed.getIOConfig().getBaseSequenceName(), + ((KinesisIndexTask) aNewTaskCapture.getValue()).getIOConfig().getBaseSequenceName() + ); + + // check that failed tasks are recreated with the same minimumMessageTime as the task it replaced, even if that + // task came from another supervisor + Assert.assertEquals( + now, + ((KinesisIndexTask) aNewTaskCapture.getValue()).getIOConfig().getMinimumMessageTime().get() + ); + Assert.assertEquals( + maxi, + ((KinesisIndexTask) aNewTaskCapture.getValue()).getIOConfig().getMaximumMessageTime().get() + ); + } + + @Test + public void testQueueNextTasksOnSuccess() throws Exception + { + supervisor = getSupervisor(2, 2, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(anyObject())).andReturn("100").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + + EasyMock.reset(taskStorage); + EasyMock.reset(taskClient); + + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of( + shardId1, + "0" + )); + // there would be 4 tasks, 2 for each task group + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(2); + + + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.replay(taskStorage); + EasyMock.replay(taskClient); + + supervisor.runInternal(); + verifyAll(); + + // test that a task succeeding causes a new task to be re-queued with the next stream range and causes any replica + // tasks to be shutdown + Capture newTasksCapture = Capture.newInstance(CaptureType.ALL); + Capture shutdownTaskIdCapture = Capture.newInstance(); + List imStillRunning = tasks.subList(1, 4); + KinesisIndexTask iAmSuccess = (KinesisIndexTask) tasks.get(0); + EasyMock.reset(taskStorage); + EasyMock.reset(taskQueue); + EasyMock.reset(taskClient); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(imStillRunning).anyTimes(); + for (Task task : imStillRunning) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.expect(taskStorage.getStatus(iAmSuccess.getId())) + .andReturn(Optional.of(TaskStatus.success(iAmSuccess.getId()))); + EasyMock.expect(taskStorage.getTask(iAmSuccess.getId())).andReturn(Optional.of((Task) iAmSuccess)).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(newTasksCapture))).andReturn(true).times(2); + EasyMock.expect(taskClient.stopAsync(EasyMock.capture(shutdownTaskIdCapture), EasyMock.eq(false))) + .andReturn(Futures.immediateFuture(true)); + EasyMock.replay(taskStorage); + EasyMock.replay(taskQueue); + EasyMock.replay(taskClient); + + supervisor.runInternal(); + verifyAll(); + + // make sure we killed the right task (sequenceName for replicas are the same) + Assert.assertTrue(shutdownTaskIdCapture.getValue().contains(iAmSuccess.getIOConfig().getBaseSequenceName())); + } + + @Test + public void testBeginPublishAndQueueNextTasks() throws Exception + { + final TaskLocation location = new TaskLocation("testHost", 1234, -1); + + supervisor = getSupervisor(2, 2, true, "PT1M", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + Collection workItems = new ArrayList<>(); + for (Task task : tasks) { + workItems.add(new TestTaskRunnerWorkItem(task, null, location)); + } + + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); + captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-1"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .times(2); + EasyMock.expect(taskClient.pauseAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "1", + shardId0, + "0" + ))) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + ))); + EasyMock.expect( + taskClient.setEndOffsetsAsync( + EasyMock.contains("sequenceName-0"), + EasyMock.eq(ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )), + EasyMock.eq(true) + ) + ).andReturn(Futures.immediateFuture(true)).times(2); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of( + shardId1, + "0" + )); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of( + shardId0, + "0" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(2); + + EasyMock.replay(taskStorage, taskRunner, taskClient, taskQueue); + + supervisor.runInternal(); + verifyAll(); + + for (Task task : captured.getValues()) { + KinesisIndexTask KinesisIndexTask = (KinesisIndexTask) task; + Assert.assertEquals(dataSchema, KinesisIndexTask.getDataSchema()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), KinesisIndexTask.getTuningConfig()); + + KinesisIndexTaskIOConfig taskConfig = KinesisIndexTask.getIOConfig(); + Assert.assertEquals("sequenceName-1", taskConfig.getBaseSequenceName()); + Assert.assertTrue("isUseTransaction", taskConfig.isUseTransaction()); + + Assert.assertEquals(stream, taskConfig.getStartPartitions().getStream()); + Assert.assertEquals( + "3", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + "1", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + } + + @Test + public void testDiscoverExistingPublishingTask() throws Exception + { + final TaskLocation location = new TaskLocation("testHost", 1234, -1); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Task task = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(task, null, location)); + + Capture captured = Capture.newInstance(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getCurrentOffsetsAsync("id1", false)) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ))); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true); + + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )); + EasyMock.expect(taskClient.getCheckpoints(EasyMock.anyString(), EasyMock.anyBoolean())) + .andReturn(checkpoints) + .anyTimes(); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + supervisor.updateCurrentAndLatestOffsets().run(); + SupervisorReport report = supervisor.getStatus(); + verifyAll(); + + Assert.assertEquals(DATASOURCE, report.getId()); + + KinesisSupervisorReportPayload payload = report.getPayload(); + + Assert.assertEquals(DATASOURCE, payload.getDataSource()); + Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); + Assert.assertEquals(2, (int) payload.getPartitions()); + Assert.assertEquals(1, (int) payload.getReplicas()); + Assert.assertEquals(stream, payload.getStream()); + Assert.assertEquals(0, payload.getActiveTasks().size()); + Assert.assertEquals(1, payload.getPublishingTasks().size()); + + TaskReportData publishingReport = payload.getPublishingTasks().get(0); + + Assert.assertEquals("id1", publishingReport.getId()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + ), publishingReport.getStartingOffsets()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ), publishingReport.getCurrentOffsets()); + + KinesisIndexTask capturedTask = captured.getValue(); + Assert.assertEquals(dataSchema, capturedTask.getDataSchema()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), capturedTask.getTuningConfig()); + + KinesisIndexTaskIOConfig capturedTaskConfig = capturedTask.getIOConfig(); + Assert.assertEquals("awsEndpoint", capturedTaskConfig.getEndpoint()); + Assert.assertEquals("sequenceName-0", capturedTaskConfig.getBaseSequenceName()); + Assert.assertTrue("isUseTransaction", capturedTaskConfig.isUseTransaction()); + + // check that the new task was created with starting sequences matching where the publishing task finished + Assert.assertEquals(stream, capturedTaskConfig.getStartPartitions().getStream()); + Assert.assertEquals( + "2", + capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + "1", + capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + + Assert.assertEquals(stream, capturedTaskConfig.getEndPartitions().getStream()); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + + @Test + public void testDiscoverExistingPublishingTaskWithDifferentPartitionAllocation() throws Exception + { + final TaskLocation location = new TaskLocation("testHost", 1234, -1); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Task task = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(task, null, location)); + + Capture captured = Capture.newInstance(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(task)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(task)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getCurrentOffsetsAsync("id1", false)) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ))); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + supervisor.updateCurrentAndLatestOffsets().run(); + SupervisorReport report = supervisor.getStatus(); + verifyAll(); + + Assert.assertEquals(DATASOURCE, report.getId()); + + KinesisSupervisorReportPayload payload = report.getPayload(); + + Assert.assertEquals(DATASOURCE, payload.getDataSource()); + Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); + Assert.assertEquals(2, (int) payload.getPartitions()); + Assert.assertEquals(1, (int) payload.getReplicas()); + Assert.assertEquals(stream, payload.getStream()); + Assert.assertEquals(0, payload.getActiveTasks().size()); + Assert.assertEquals(1, payload.getPublishingTasks().size()); + + TaskReportData publishingReport = payload.getPublishingTasks().get(0); + + Assert.assertEquals("id1", publishingReport.getId()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + ), publishingReport.getStartingOffsets()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ), publishingReport.getCurrentOffsets()); + + KinesisIndexTask capturedTask = captured.getValue(); + Assert.assertEquals(dataSchema, capturedTask.getDataSchema()); + Assert.assertEquals(tuningConfig.convertToTaskTuningConfig(), capturedTask.getTuningConfig()); + + KinesisIndexTaskIOConfig capturedTaskConfig = capturedTask.getIOConfig(); + Assert.assertEquals("awsEndpoint", capturedTaskConfig.getEndpoint()); + Assert.assertEquals("sequenceName-0", capturedTaskConfig.getBaseSequenceName()); + Assert.assertTrue("isUseTransaction", capturedTaskConfig.isUseTransaction()); + + // check that the new task was created with starting sequences matching where the publishing task finished + Assert.assertEquals(stream, capturedTaskConfig.getStartPartitions().getStream()); + Assert.assertEquals( + "2", + capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + "1", + capturedTaskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + + Assert.assertEquals(stream, capturedTaskConfig.getEndPartitions().getStream()); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertEquals( + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + capturedTaskConfig.getEndPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + + @Test + public void testDiscoverExistingPublishingAndReadingTask() throws Exception + { + final TaskLocation location1 = new TaskLocation("testHost", 1234, -1); + final TaskLocation location2 = new TaskLocation("testHost2", 145, -1); + final DateTime startTime = DateTimes.nowUtc(); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>( + "stream", + ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + ) + ), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getCurrentOffsetsAsync("id1", false)) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ))); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )); + EasyMock.expect(taskClient.getCurrentOffsetsAsync("id2", false)) + .andReturn(Futures.immediateFuture((Map) ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ))); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + + // since id1 is publishing, so getCheckpoints wouldn't be called for it + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + supervisor.updateCurrentAndLatestOffsets().run(); + SupervisorReport report = supervisor.getStatus(); + verifyAll(); + + Assert.assertEquals(DATASOURCE, report.getId()); + + KinesisSupervisorReportPayload payload = report.getPayload(); + + Assert.assertEquals(DATASOURCE, payload.getDataSource()); + Assert.assertEquals(3600L, (long) payload.getDurationSeconds()); + Assert.assertEquals(2, (int) payload.getPartitions()); + Assert.assertEquals(1, (int) payload.getReplicas()); + Assert.assertEquals(stream, payload.getStream()); + Assert.assertEquals(1, payload.getActiveTasks().size()); + Assert.assertEquals(1, payload.getPublishingTasks().size()); + + TaskReportData activeReport = payload.getActiveTasks().get(0); + TaskReportData publishingReport = payload.getPublishingTasks().get(0); + + Assert.assertEquals("id2", activeReport.getId()); + Assert.assertEquals(startTime, activeReport.getStartTime()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ), activeReport.getStartingOffsets()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ), activeReport.getCurrentOffsets()); + + Assert.assertEquals("id1", publishingReport.getId()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + ), publishingReport.getStartingOffsets()); + Assert.assertEquals(ImmutableMap.of( + shardId1, + "2", + shardId0, + "1" + ), publishingReport.getCurrentOffsets()); + } + + @Test + public void testKillUnresponsiveTasksWhileGettingStartTime() throws Exception + { + supervisor = getSupervisor(2, 2, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + + EasyMock.reset(taskStorage, taskClient, taskQueue); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of( + shardId1, + + "0", + + shardId0, + + "0" + )); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of(shardId0, "0")); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(2); + + + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(task.getId())) + .andReturn(Futures.immediateFuture(Status.NOT_STARTED)); + EasyMock.expect(taskClient.getStartTimeAsync(task.getId())) + .andReturn(Futures.immediateFailedFuture(new RuntimeException())); + taskQueue.shutdown(task.getId(), "Task [%s] failed to return start time, killing task", task.getId()); + } + EasyMock.replay(taskStorage, taskClient, taskQueue); + + supervisor.runInternal(); + verifyAll(); + } + + @Test + public void testKillUnresponsiveTasksWhilePausing() throws Exception + { + final TaskLocation location = new TaskLocation("testHost", 1234, -1); + + supervisor = getSupervisor(2, 2, true, "PT1M", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + Collection workItems = new ArrayList<>(); + for (Task task : tasks) { + workItems.add(new TestTaskRunnerWorkItem(task, null, location)); + } + + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of( + shardId1, + "0" + )); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of(shardId0, "0")); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(2); + + captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-1"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .times(2); + EasyMock.expect(taskClient.pauseAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFailedFuture(new RuntimeException())).times(2); + taskQueue.shutdown( + EasyMock.contains("sequenceName-0"), + EasyMock.eq("An exception occured while waiting for task [%s] to pause: [%s]"), + EasyMock.contains("sequenceName-0"), + EasyMock.anyString() + ); + EasyMock.expectLastCall().times(2); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + + EasyMock.replay(taskStorage, taskRunner, taskClient, taskQueue); + + supervisor.runInternal(); + verifyAll(); + + for (Task task : captured.getValues()) { + KinesisIndexTaskIOConfig taskConfig = ((KinesisIndexTask) task).getIOConfig(); + Assert.assertEquals( + "0", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + Assert.assertNull( + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId0) + ); + } + } + + @Test + public void testKillUnresponsiveTasksWhileSettingEndOffsets() throws Exception + { + final TaskLocation location = new TaskLocation("testHost", 1234, -1); + + supervisor = getSupervisor(2, 2, true, "PT1M", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Capture captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(4); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + List tasks = captured.getValues(); + Collection workItems = new ArrayList<>(); + for (Task task : tasks) { + workItems.add(new TestTaskRunnerWorkItem(task, null, location)); + } + + EasyMock.reset(taskStorage, taskRunner, taskClient, taskQueue); + + TreeMap> checkpoints1 = new TreeMap<>(); + checkpoints1.put(0, ImmutableMap.of( + shardId1, + "0" + )); + TreeMap> checkpoints2 = new TreeMap<>(); + checkpoints2.put(0, ImmutableMap.of(shardId0, "0")); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-0"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints1)) + .times(2); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("sequenceName-1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints2)) + .times(2); + + captured = Capture.newInstance(CaptureType.ALL); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(tasks).anyTimes(); + for (Task task : tasks) { + EasyMock.expect(taskStorage.getStatus(task.getId())) + .andReturn(Optional.of(TaskStatus.running(task.getId()))) + .anyTimes(); + EasyMock.expect(taskStorage.getTask(task.getId())).andReturn(Optional.of(task)).anyTimes(); + } + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) + .anyTimes(); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc().minusMinutes(2))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())); + EasyMock.expect(taskClient.getStartTimeAsync(EasyMock.contains("sequenceName-1"))) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .times(2); + EasyMock.expect(taskClient.pauseAsync(EasyMock.contains("sequenceName-0"))) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "1" + ))) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "3" + ))); + EasyMock.expect( + taskClient.setEndOffsetsAsync( + EasyMock.contains("sequenceName-0"), + EasyMock.eq(ImmutableMap.of( + shardId1, + "3" + )), + EasyMock.eq(true) + ) + ).andReturn(Futures.immediateFailedFuture(new RuntimeException())).times(2); + taskQueue.shutdown( + EasyMock.contains("sequenceName-0"), + EasyMock.eq("All tasks in group [%s] failed to transition to publishing state"), + EasyMock.eq(0) + ); + EasyMock.expectLastCall().times(2); + EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2); + + EasyMock.replay(taskStorage, taskRunner, taskClient, taskQueue); + + supervisor.runInternal(); + verifyAll(); + + for (Task task : captured.getValues()) { + KinesisIndexTaskIOConfig taskConfig = ((KinesisIndexTask) task).getIOConfig(); + Assert.assertEquals( + "0", + taskConfig.getStartPartitions().getPartitionSequenceNumberMap().get(shardId1) + ); + } + } + + @Test(expected = IllegalStateException.class) + public void testStopNotStarted() + { + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + supervisor.stop(false); + } + + @Test + public void testStop() + { + supervisorRecordSupplier.close(); + expectLastCall().anyTimes(); + + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + taskClient.close(); + taskRunner.unregisterListener(StringUtils.format("KinesisSupervisor-%s", DATASOURCE)); + replayAll(); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + supervisor.start(); + supervisor.stop(false); + + verifyAll(); + } + + @Test + public void testStopGracefully() throws Exception + { + final TaskLocation location1 = new TaskLocation("testHost", 1234, -1); + final TaskLocation location2 = new TaskLocation("testHost2", 145, -1); + final DateTime startTime = DateTimes.nowUtc(); + + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStatusAsync("id3")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )); + + // getCheckpoints will not be called for id1 as it is in publishing state + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + EasyMock.reset(taskRunner, taskClient, taskQueue); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskClient.pauseAsync("id2")) + .andReturn(Futures.immediateFuture((Map) ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ))); + EasyMock.expect(taskClient.setEndOffsetsAsync("id2", ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ), true)) + .andReturn(Futures.immediateFuture(true)); + taskQueue.shutdown("id3", "Killing task for graceful shutdown"); + expectLastCall().times(1); + taskQueue.shutdown("id3", "Killing task [%s] which hasn't been assigned to a worker", "id3"); + expectLastCall().times(1); + + EasyMock.replay(taskRunner, taskClient, taskQueue); + + supervisor.gracefulShutdownInternal(); + verifyAll(); + } + + @Test + public void testResetNoTasks() throws Exception + { + expect(supervisorRecordSupplier.getPartitionIds(anyObject())).andReturn(Collections.emptySet()).anyTimes(); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + EasyMock.reset(indexerMetadataStorageCoordinator); + EasyMock.expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); + EasyMock.replay(indexerMetadataStorageCoordinator); + + supervisor.resetInternal(null); + verifyAll(); + + } + + @Test + public void testResetDataSourceMetadata() throws Exception + { + expect(supervisorRecordSupplier.getPartitionIds(anyObject())).andReturn(Collections.emptySet()).anyTimes(); + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + Capture captureDataSource = EasyMock.newCapture(); + Capture captureDataSourceMetadata = EasyMock.newCapture(); + + KinesisDataSourceMetadata KinesisDataSourceMetadata = new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + ) + )); + + KinesisDataSourceMetadata resetMetadata = new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + ) + )); + + KinesisDataSourceMetadata expectedMetadata = new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + ) + )); + + EasyMock.reset(indexerMetadataStorageCoordinator); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)) + .andReturn(KinesisDataSourceMetadata); + EasyMock.expect(indexerMetadataStorageCoordinator.resetDataSourceMetadata( + EasyMock.capture(captureDataSource), + EasyMock.capture(captureDataSourceMetadata) + )).andReturn(true); + EasyMock.replay(indexerMetadataStorageCoordinator); + + try { + supervisor.resetInternal(resetMetadata); + } + catch (NullPointerException npe) { + // Expected as there will be an attempt to EasyMock.reset partitionGroups sequences to NOT_SET + // however there would be no entries in the map as we have not put nay data in kafka + Assert.assertTrue(npe.getCause() == null); + } + verifyAll(); + + Assert.assertEquals(captureDataSource.getValue(), DATASOURCE); + Assert.assertEquals(captureDataSourceMetadata.getValue(), expectedMetadata); + } + + @Test + public void testResetNoDataSourceMetadata() throws Exception + { + expect(supervisorRecordSupplier.getPartitionIds(anyObject())).andReturn(Collections.emptySet()).anyTimes(); + supervisor = getSupervisor(1, 1, true, "PT1H", null, null); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + KinesisDataSourceMetadata resetMetadata = new KinesisDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + ImmutableMap.of( + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + ) + )); + + EasyMock.reset(indexerMetadataStorageCoordinator); + // no DataSourceMetadata in metadata store + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(null); + EasyMock.replay(indexerMetadataStorageCoordinator); + + supervisor.resetInternal(resetMetadata); + verifyAll(); + } + + @Test + public void testResetRunningTasks() throws Exception + { + final TaskLocation location1 = new TaskLocation("testHost", 1234, -1); + final TaskLocation location2 = new TaskLocation("testHost2", 145, -1); + final DateTime startTime = DateTimes.nowUtc(); + + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + + "0", + + shardId0, + + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + + "3", + + shardId0, + + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + + "3", + + shardId0, + + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStatusAsync("id3")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + + "3", + + shardId0, + + "1" + )); + + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + + "3", + + shardId0, + + "1" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + EasyMock.reset(taskQueue, indexerMetadataStorageCoordinator); + EasyMock.expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); + taskQueue.shutdown("id2", "DataSourceMetadata is not found while reset"); + taskQueue.shutdown("id3", "DataSourceMetadata is not found while reset"); + EasyMock.replay(taskQueue, indexerMetadataStorageCoordinator); + + supervisor.resetInternal(null); + verifyAll(); + } + + @Test + public void testNoDataIngestionTasks() throws Exception + { + final DateTime startTime = DateTimes.nowUtc(); + supervisor = getSupervisor(2, 1, true, "PT1H", null, null); + + //not adding any events + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStatusAsync("id3")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); + + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + EasyMock.reset(taskQueue, indexerMetadataStorageCoordinator); + EasyMock.expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); + taskQueue.shutdown("id1", "DataSourceMetadata is not found while reset"); + taskQueue.shutdown("id2", "DataSourceMetadata is not found while reset"); + taskQueue.shutdown("id3", "DataSourceMetadata is not found while reset"); + EasyMock.replay(taskQueue, indexerMetadataStorageCoordinator); + + supervisor.resetInternal(null); + verifyAll(); + } + + + @Test(timeout = 60_000L) + public void testCheckpointForInactiveTaskGroup() + throws InterruptedException, ExecutionException, TimeoutException, JsonProcessingException + { + supervisor = getSupervisor(2, 1, true, "PT1S", null, null, false); + //not adding any events + final Task id1; + id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final TaskLocation location1 = new TaskLocation("testHost", 1234, -1); + final TaskLocation location2 = new TaskLocation("testHost2", 145, -1); + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + expect( + indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata( + null) + ).anyTimes(); + expect(taskClient.getStatusAsync("id1")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id2")).andReturn(Futures.immediateFuture(Status.READING)); + expect(taskClient.getStatusAsync("id3")).andReturn(Futures.immediateFuture(Status.READING)); + + final DateTime startTime = DateTimes.nowUtc(); + expect(taskClient.getStartTimeAsync("id1")).andReturn(Futures.immediateFuture(startTime)); + expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); + + final TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )); + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id1"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + + final Map fakeCheckpoints = Collections.emptyMap(); + supervisor.moveTaskGroupToPendingCompletion(0); + supervisor.checkpoint( + 0, + ((KinesisIndexTask) id1).getIOConfig().getBaseSequenceName(), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, checkpoints.get(0))), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, fakeCheckpoints)) + ); + + while (supervisor.getNoticesQueueSize() > 0) { + Thread.sleep(100); + } + + verifyAll(); + + Assert.assertNull(serviceEmitter.getStackTrace(), serviceEmitter.getStackTrace()); + Assert.assertNull(serviceEmitter.getExceptionMessage(), serviceEmitter.getExceptionMessage()); + Assert.assertNull(serviceEmitter.getExceptionClass()); + } + + @Test(timeout = 60_000L) + public void testCheckpointForUnknownTaskGroup() + throws InterruptedException + { + supervisor = getSupervisor(2, 1, true, "PT1S", null, null, false); + + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + //not adding any events + final Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "10", + shardId0, + "20" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + expect( + indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata( + null) + ).anyTimes(); + + replayAll(); + + supervisor.start(); + + supervisor.checkpoint( + 0, + ((KinesisIndexTask) id1).getIOConfig().getBaseSequenceName(), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, Collections.emptyMap())), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, Collections.emptyMap())) + ); + + while (supervisor.getNoticesQueueSize() > 0) { + Thread.sleep(100); + } + + verifyAll(); + + while (serviceEmitter.getStackTrace() == null) { + Thread.sleep(100); + } + + Assert.assertTrue(serviceEmitter.getStackTrace() + .startsWith("org.apache.druid.java.util.common.ISE: WTH?! cannot find")); + Assert.assertEquals( + "WTH?! cannot find taskGroup [0] among all activelyReadingTaskGroups [{}]", + serviceEmitter.getExceptionMessage() + ); + Assert.assertEquals(ISE.class, serviceEmitter.getExceptionClass()); + } + + @Test(timeout = 60_000L) + public void testCheckpointWithNullTaskGroupId() + throws InterruptedException, ExecutionException, TimeoutException, JsonProcessingException + { + supervisor = getSupervisor(1, 3, true, "PT1S", null, null, false); + //not adding any events + final Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + final Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(Collections.emptySet()).anyTimes(); + expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + expect( + indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn(new KinesisDataSourceMetadata( + null) + ).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + expect(taskClient.getStatusAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(Status.READING)) + .anyTimes(); + final TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of(shardId1, "0")); + expect(taskClient.getCheckpointsAsync(EasyMock.anyString(), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(3); + expect(taskClient.getStartTimeAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(DateTimes.nowUtc())) + .anyTimes(); + expect(taskClient.pauseAsync(EasyMock.anyString())) + .andReturn(Futures.immediateFuture(ImmutableMap.of(shardId1, "10"))) + .anyTimes(); + expect(taskClient.setEndOffsetsAsync( + EasyMock.anyString(), + EasyMock.eq(ImmutableMap.of("0", "10")), + EasyMock.anyBoolean() + )) + .andReturn(Futures.immediateFuture(true)) + .anyTimes(); + + replayAll(); + + supervisor.start(); + + supervisor.runInternal(); + + final TreeMap> newCheckpoints = new TreeMap<>(); + newCheckpoints.put(0, ImmutableMap.of(shardId1, "10")); + supervisor.checkpoint( + null, + ((KinesisIndexTask) id1).getIOConfig().getBaseSequenceName(), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, checkpoints.get(0))), + new KinesisDataSourceMetadata(new SeekableStreamPartitions<>(stream, newCheckpoints.get(0))) + ); + + while (supervisor.getNoticesQueueSize() > 0) { + Thread.sleep(100); + } + + verifyAll(); + } + + + @Test + public void testSuspendedNoRunningTasks() throws Exception + { + supervisor = getSupervisor(1, 1, true, "PT1H", null, null, true); + + expect(supervisorRecordSupplier.getPartitionIds(anyObject())).andReturn(Collections.emptySet()).anyTimes(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + // this asserts that taskQueue.add does not in fact get called because supervisor should be suspended + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andAnswer((IAnswer) () -> { + Assert.fail(); + return null; + }).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + } + + @Test + public void testSuspendedRunningTasks() throws Exception + { + // graceful shutdown is expected to be called on running tasks since state is suspended + + final TaskLocation location1 = new TaskLocation("testHost", 1234, -1); + final TaskLocation location2 = new TaskLocation("testHost2", 145, -1); + final DateTime startTime = DateTimes.nowUtc(); + + supervisor = getSupervisor(2, 1, true, "PT1H", null, null, true); + + supervisorRecordSupplier.assign(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getPartitionIds(stream)).andReturn(ImmutableSet.of(shardId1, shardId0)).anyTimes(); + expect(supervisorRecordSupplier.getAssignment()).andReturn(ImmutableSet.of(shard1Partition, shard0Partition)) + .anyTimes(); + supervisorRecordSupplier.seekToLatest(anyObject()); + expectLastCall().anyTimes(); + expect(supervisorRecordSupplier.getEarliestSequenceNumber(anyObject())).andReturn("0").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard1Partition)).andReturn("12").anyTimes(); + expect(supervisorRecordSupplier.getLatestSequenceNumber(shard0Partition)).andReturn("1").anyTimes(); + supervisorRecordSupplier.seek(anyObject(), anyString()); + expectLastCall().anyTimes(); + + + Task id1 = createKinesisIndexTask( + "id1", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "0", + shardId0, + "0" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id2 = createKinesisIndexTask( + "id2", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Task id3 = createKinesisIndexTask( + "id3", + DATASOURCE, + 0, + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )), + new SeekableStreamPartitions<>("stream", ImmutableMap.of( + shardId1, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER, + shardId0, + SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER + )), + null, + null + ); + + Collection workItems = new ArrayList<>(); + workItems.add(new TestTaskRunnerWorkItem(id1, null, location1)); + workItems.add(new TestTaskRunnerWorkItem(id2, null, location2)); + + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(workItems).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of(id1, id2, id3)).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id1")).andReturn(Optional.of(TaskStatus.running("id1"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id2")).andReturn(Optional.of(TaskStatus.running("id2"))).anyTimes(); + EasyMock.expect(taskStorage.getStatus("id3")).andReturn(Optional.of(TaskStatus.running("id3"))).anyTimes(); + EasyMock.expect(taskStorage.getTask("id1")).andReturn(Optional.of(id1)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id2")).andReturn(Optional.of(id2)).anyTimes(); + EasyMock.expect(taskStorage.getTask("id3")).andReturn(Optional.of(id3)).anyTimes(); + EasyMock.expect(indexerMetadataStorageCoordinator.getDataSourceMetadata(DATASOURCE)).andReturn( + new KinesisDataSourceMetadata( + null + ) + ).anyTimes(); + EasyMock.expect(taskClient.getStatusAsync("id1")) + .andReturn(Futures.immediateFuture(Status.PUBLISHING)); + EasyMock.expect(taskClient.getStatusAsync("id2")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStatusAsync("id3")) + .andReturn(Futures.immediateFuture(Status.READING)); + EasyMock.expect(taskClient.getStartTimeAsync("id2")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getStartTimeAsync("id3")).andReturn(Futures.immediateFuture(startTime)); + EasyMock.expect(taskClient.getEndOffsets("id1")).andReturn(ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )); + + // getCheckpoints will not be called for id1 as it is in publishing state + TreeMap> checkpoints = new TreeMap<>(); + checkpoints.put(0, ImmutableMap.of( + shardId1, + "3", + shardId0, + "1" + )); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id2"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.contains("id3"), EasyMock.anyBoolean())) + .andReturn(Futures.immediateFuture(checkpoints)) + .times(1); + + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + + EasyMock.expect(taskClient.pauseAsync("id2")) + .andReturn(Futures.immediateFuture(ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ))); + EasyMock.expect(taskClient.setEndOffsetsAsync("id2", ImmutableMap.of( + shardId1, + "12", + shardId0, + "1" + ), true)) + .andReturn(Futures.immediateFuture(true)); + taskQueue.shutdown("id3", "Killing task for graceful shutdown"); + expectLastCall().times(1); + taskQueue.shutdown("id3", "Killing task [%s] which hasn't been assigned to a worker", "id3"); + expectLastCall().times(1); + + replayAll(); + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + } + + @Test + public void testResetSuspended() throws Exception + { + expect(supervisorRecordSupplier.getPartitionIds(anyObject())).andReturn(Collections.emptySet()).anyTimes(); + EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes(); + EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes(); + EasyMock.expect(taskRunner.getRunningTasks()).andReturn(Collections.EMPTY_LIST).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasks()).andReturn(ImmutableList.of()).anyTimes(); + taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class), EasyMock.anyObject(Executor.class)); + replayAll(); + + supervisor = getSupervisor(1, 1, true, "PT1H", null, null, true); + supervisor.start(); + supervisor.runInternal(); + verifyAll(); + + EasyMock.reset(indexerMetadataStorageCoordinator); + EasyMock.expect(indexerMetadataStorageCoordinator.deleteDataSourceMetadata(DATASOURCE)).andReturn(true); + EasyMock.replay(indexerMetadataStorageCoordinator); + + supervisor.resetInternal(null); + verifyAll(); + } + + + @Test + public void testGetCurrentTotalStats() + { + supervisor = getSupervisor(1, 2, true, "PT1H", null, null, false); + supervisor.addTaskGroupToActivelyReadingTaskGroup( + supervisor.getTaskGroupIdForPartition("0"), + ImmutableMap.of("0", "0"), + Optional.absent(), + Optional.absent(), + ImmutableSet.of("task1"), + ImmutableSet.of() + ); + + supervisor.addTaskGroupToPendingCompletionTaskGroup( + supervisor.getTaskGroupIdForPartition("1"), + ImmutableMap.of("0", "0"), + Optional.absent(), + Optional.absent(), + ImmutableSet.of("task2"), + ImmutableSet.of() + ); + + expect(taskClient.getMovingAveragesAsync("task1")).andReturn(Futures.immediateFuture(ImmutableMap.of( + "prop1", + "val1" + ))).times(1); + + expect(taskClient.getMovingAveragesAsync("task2")).andReturn(Futures.immediateFuture(ImmutableMap.of( + "prop2", + "val2" + ))).times(1); + + replayAll(); + + Map> stats = supervisor.getStats(); + + verifyAll(); + + Assert.assertEquals(2, stats.size()); + Assert.assertEquals(ImmutableSet.of("0", "1"), stats.keySet()); + Assert.assertEquals(ImmutableMap.of("task1", ImmutableMap.of("prop1", "val1")), stats.get("0")); + Assert.assertEquals(ImmutableMap.of("task2", ImmutableMap.of("prop2", "val2")), stats.get("1")); + } + + private KinesisSupervisor getSupervisor( + int replicas, + int taskCount, + boolean useEarliestOffset, + String duration, + Period lateMessageRejectionPeriod, + Period earlyMessageRejectionPeriod + ) + { + return getSupervisor( + replicas, + taskCount, + useEarliestOffset, + duration, + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod, + false, + null, + null + ); + } + + private KinesisSupervisor getSupervisor( + int replicas, + int taskCount, + boolean useEarliestOffset, + String duration, + Period lateMessageRejectionPeriod, + Period earlyMessageRejectionPeriod, + boolean suspended, + Integer recordsPerFetch, + Integer fetchDelayMillis + + ) + { + KinesisSupervisorIOConfig KinesisSupervisorIOConfig = new KinesisSupervisorIOConfig( + stream, + "awsEndpoint", + null, + replicas, + taskCount, + new Period(duration), + new Period("P1D"), + new Period("PT30S"), + useEarliestOffset, + new Period("PT30M"), + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod, + recordsPerFetch, + fetchDelayMillis, + null, + null, + false + ); + + KinesisIndexTaskClientFactory taskClientFactory = new KinesisIndexTaskClientFactory( + null, + null + ) + { + @Override + public KinesisIndexTaskClient build( + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ) + { + Assert.assertEquals(TEST_CHAT_THREADS, numThreads); + Assert.assertEquals(TEST_HTTP_TIMEOUT.toStandardDuration(), httpTimeout); + Assert.assertEquals(TEST_CHAT_RETRIES, numRetries); + return taskClient; + } + }; + + return new TestableKinesisSupervisor( + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + objectMapper, + new KinesisSupervisorSpec( + dataSchema, + tuningConfig, + KinesisSupervisorIOConfig, + null, + suspended, + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + objectMapper, + new NoopServiceEmitter(), + new DruidMonitorSchedulerConfig(), + rowIngestionMetersFactory, + null + ), + rowIngestionMetersFactory + ); + } + + private static DataSchema getDataSchema(String dataSource) + { + List dimensions = new ArrayList<>(); + dimensions.add(StringDimensionSchema.create("dim1")); + dimensions.add(StringDimensionSchema.create("dim2")); + + return new DataSchema( + dataSource, + objectMapper.convertValue( + new StringInputRowParser( + new JSONParseSpec( + new TimestampSpec("timestamp", "iso", null), + new DimensionsSpec( + dimensions, + null, + null + ), + new JSONPathSpec(true, ImmutableList.of()), + ImmutableMap.of() + ), + StandardCharsets.UTF_8.name() + ), + Map.class + ), + new AggregatorFactory[]{new CountAggregatorFactory("rows")}, + new UniformGranularitySpec( + Granularities.HOUR, + Granularities.NONE, + ImmutableList.of() + ), + null, + objectMapper + ); + } + + private static List JB( + String timestamp, + String dim1, + String dim2, + String dimLong, + String dimFloat, + String met1 + ) + { + try { + return Collections.singletonList(new ObjectMapper().writeValueAsBytes( + ImmutableMap.builder() + .put("timestamp", timestamp) + .put("dim1", dim1) + .put("dim2", dim2) + .put("dimLong", dimLong) + .put("dimFloat", dimFloat) + .put("met1", met1) + .build() + )); + } + catch (Exception e) { + throw Throwables.propagate(e); + } + } + + private KinesisIndexTask createKinesisIndexTask( + String id, + String dataSource, + int taskGroupId, + SeekableStreamPartitions startPartitions, + SeekableStreamPartitions endPartitions, + DateTime minimumMessageTime, + DateTime maximumMessageTime + ) + { + return new KinesisIndexTask( + id, + null, + getDataSchema(dataSource), + tuningConfig, + new KinesisIndexTaskIOConfig( + null, + "sequenceName-" + taskGroupId, + startPartitions, + endPartitions, + true, + minimumMessageTime, + maximumMessageTime, + "awsEndpoint", + null, + null, + null, + null, + null, + false + ), + Collections.emptyMap(), + null, + null, + rowIngestionMetersFactory, + null + ); + } + + private KinesisSupervisor getSupervisor( + int replicas, + int taskCount, + boolean useEarliestOffset, + String duration, + Period lateMessageRejectionPeriod, + Period earlyMessageRejectionPeriod, + boolean suspended + ) + { + KinesisSupervisorIOConfig KinesisSupervisorIOConfig = new KinesisSupervisorIOConfig( + stream, + "awsEndpoint", + null, + replicas, + taskCount, + new Period(duration), + new Period("P1D"), + new Period("PT30S"), + useEarliestOffset, + new Period("PT30M"), + lateMessageRejectionPeriod, + earlyMessageRejectionPeriod, + null, + null, + null, + null, + false + ); + + KinesisIndexTaskClientFactory taskClientFactory = new KinesisIndexTaskClientFactory( + null, + null + ) + { + @Override + public KinesisIndexTaskClient build( + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ) + { + Assert.assertEquals(TEST_CHAT_THREADS, numThreads); + Assert.assertEquals(TEST_HTTP_TIMEOUT.toStandardDuration(), httpTimeout); + Assert.assertEquals(TEST_CHAT_RETRIES, numRetries); + return taskClient; + } + }; + + return new TestableKinesisSupervisor( + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + objectMapper, + new KinesisSupervisorSpec( + dataSchema, + tuningConfig, + KinesisSupervisorIOConfig, + null, + suspended, + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + objectMapper, + new NoopServiceEmitter(), + new DruidMonitorSchedulerConfig(), + rowIngestionMetersFactory, + null + ), + rowIngestionMetersFactory + ); + } + + private static class TestTaskRunnerWorkItem extends TaskRunnerWorkItem + { + private final String taskType; + private final TaskLocation location; + private final String dataSource; + + public TestTaskRunnerWorkItem(Task task, ListenableFuture result, TaskLocation location) + { + super(task.getId(), result); + this.taskType = task.getType(); + this.location = location; + this.dataSource = task.getDataSource(); + } + + @Override + public TaskLocation getLocation() + { + return location; + } + + @Override + public String getTaskType() + { + return taskType; + } + + @Override + public String getDataSource() + { + return dataSource; + } + + } + + private class TestableKinesisSupervisor extends KinesisSupervisor + { + private final KinesisSupervisorSpec spec; + + public TestableKinesisSupervisor( + TaskStorage taskStorage, + TaskMaster taskMaster, + IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + KinesisIndexTaskClientFactory taskClientFactory, + ObjectMapper mapper, + KinesisSupervisorSpec spec, + RowIngestionMetersFactory rowIngestionMetersFactory + ) + { + super( + taskStorage, + taskMaster, + indexerMetadataStorageCoordinator, + taskClientFactory, + mapper, + spec, + rowIngestionMetersFactory, + null + ); + this.spec = spec; + } + + @Override + protected String generateSequenceName( + Map startPartitions, + Optional minimumMessageTime, + Optional maximumMessageTime + ) + { + final int groupId = getTaskGroupIdForPartition(startPartitions.keySet().iterator().next()); + return StringUtils.format("sequenceName-%d", groupId); + } + + @Override + protected RecordSupplier setupRecordSupplier() + { + return supervisorRecordSupplier; + } + + + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfigTest.java b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfigTest.java new file mode 100644 index 000000000000..370f3ea2ea0d --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTuningConfigTest.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.kinesis.supervisor; + +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.indexing.kinesis.KinesisIndexingServiceModule; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.indexing.TuningConfig; +import org.joda.time.Duration; +import org.joda.time.Period; +import org.junit.Assert; +import org.junit.Test; + +import java.io.File; + +public class KinesisSupervisorTuningConfigTest +{ + private final ObjectMapper mapper; + + public KinesisSupervisorTuningConfigTest() + { + mapper = new DefaultObjectMapper(); + mapper.registerModules((Iterable) new KinesisIndexingServiceModule().getJacksonModules()); + } + + @Test + public void testSerdeWithDefaults() throws Exception + { + String jsonStr = "{\"type\": \"kinesis\"}"; + + KinesisSupervisorTuningConfig config = (KinesisSupervisorTuningConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + TuningConfig.class + ) + ), + TuningConfig.class + ); + + Assert.assertNotNull(config.getBasePersistDirectory()); + Assert.assertEquals(1000000, config.getMaxRowsInMemory()); + Assert.assertEquals(5_000_000, config.getMaxRowsPerSegment()); + Assert.assertEquals(new Period("PT10M"), config.getIntermediatePersistPeriod()); + Assert.assertEquals(0, config.getMaxPendingPersists()); + Assert.assertEquals(new IndexSpec(), config.getIndexSpec()); + Assert.assertEquals(true, config.getBuildV9Directly()); + Assert.assertEquals(false, config.isReportParseExceptions()); + Assert.assertEquals(0, config.getHandoffConditionTimeout()); + Assert.assertNull(config.getWorkerThreads()); + Assert.assertNull(config.getChatThreads()); + Assert.assertEquals(8L, (long) config.getChatRetries()); + Assert.assertEquals(Duration.standardSeconds(10), config.getHttpTimeout()); + Assert.assertEquals(Duration.standardSeconds(80), config.getShutdownTimeout()); + } + + @Test + public void testSerdeWithNonDefaults() throws Exception + { + String jsonStr = "{\n" + + " \"type\": \"kinesis\",\n" + + " \"basePersistDirectory\": \"/tmp/xxx\",\n" + + " \"maxRowsInMemory\": 100,\n" + + " \"maxRowsPerSegment\": 100,\n" + + " \"intermediatePersistPeriod\": \"PT1H\",\n" + + " \"maxPendingPersists\": 100,\n" + + " \"buildV9Directly\": false,\n" + + " \"reportParseExceptions\": true,\n" + + " \"handoffConditionTimeout\": 100,\n" + + " \"workerThreads\": 12,\n" + + " \"chatThreads\": 13,\n" + + " \"chatRetries\": 14,\n" + + " \"httpTimeout\": \"PT15S\",\n" + + " \"shutdownTimeout\": \"PT95S\"\n" + + "}"; + + KinesisSupervisorTuningConfig config = (KinesisSupervisorTuningConfig) mapper.readValue( + mapper.writeValueAsString( + mapper.readValue( + jsonStr, + TuningConfig.class + ) + ), + TuningConfig.class + ); + + Assert.assertEquals(new File("/tmp/xxx"), config.getBasePersistDirectory()); + Assert.assertEquals(100, config.getMaxRowsInMemory()); + Assert.assertEquals(100, config.getMaxRowsPerSegment()); + Assert.assertEquals(new Period("PT1H"), config.getIntermediatePersistPeriod()); + Assert.assertEquals(100, config.getMaxPendingPersists()); + Assert.assertEquals(true, config.getBuildV9Directly()); + Assert.assertEquals(true, config.isReportParseExceptions()); + Assert.assertEquals(100, config.getHandoffConditionTimeout()); + Assert.assertEquals(12, (int) config.getWorkerThreads()); + Assert.assertEquals(13, (int) config.getChatThreads()); + Assert.assertEquals(14L, (long) config.getChatRetries()); + Assert.assertEquals(Duration.standardSeconds(15), config.getHttpTimeout()); + Assert.assertEquals(Duration.standardSeconds(95), config.getShutdownTimeout()); + } +} diff --git a/extensions-core/kinesis-indexing-service/src/test/resources/log4j2.xml b/extensions-core/kinesis-indexing-service/src/test/resources/log4j2.xml new file mode 100644 index 000000000000..bca6c69fdb1b --- /dev/null +++ b/extensions-core/kinesis-indexing-service/src/test/resources/log4j2.xml @@ -0,0 +1,35 @@ + + + + + + + + + + + + + + + + + + diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/IndexTaskClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/IndexTaskClient.java index 148dc6c87045..1b65ad490c2f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/IndexTaskClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/IndexTaskClient.java @@ -21,7 +21,9 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JavaType; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.type.TypeFactory; import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Optional; @@ -55,6 +57,7 @@ import java.net.Socket; import java.net.URL; import java.nio.charset.StandardCharsets; +import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -143,6 +146,11 @@ protected RetryPolicy newRetryPolicy() return retryPolicyFactory.makeRetryPolicy(); } + protected T deserialize(String content, JavaType type) throws IOException + { + return objectMapper.readValue(content, type); + } + protected T deserialize(String content, TypeReference typeReference) throws IOException { return objectMapper.readValue(content, typeReference); @@ -153,6 +161,33 @@ protected T deserialize(String content, Class typeReference) throws IOExc return objectMapper.readValue(content, typeReference); } + protected T deserializeMap(String content, Class mapClass, Class keyClass, Class valueClass) + throws IOException + { + return deserialize(content, objectMapper.getTypeFactory().constructMapType(mapClass, keyClass, valueClass)); + } + + protected T deserializeNestedValueMap( + String content, + Class mapClass, + Class keyClass, + Class valueMapClass, + Class valueMapClassKey, + Class valueMapClassValue + ) + throws IOException + { + TypeFactory factory = objectMapper.getTypeFactory(); + return deserialize( + content, + factory.constructMapType( + mapClass, + factory.constructType(keyClass), + factory.constructMapType(valueMapClass, valueMapClassKey, valueMapClassValue) + ) + ); + } + protected byte[] serialize(Object value) throws JsonProcessingException { return objectMapper.writeValueAsBytes(value); @@ -400,14 +435,14 @@ private FullResponseHolder submitRequest(Request request) throws IOException, Ch /** * Throws if it's possible to throw the given Throwable. - * + *

* - The input throwable shouldn't be null. * - If Throwable is an {@link ExecutionException}, this calls itself recursively with the cause of ExecutionException. * - If Throwable is an {@link IOException} or a {@link ChannelException}, this simply throws it. * - If Throwable is an {@link InterruptedException}, this interrupts the current thread and throws a RuntimeException - * wrapping the InterruptedException + * wrapping the InterruptedException * - Otherwise, this simply returns the given Throwable. - * + *

* Note that if the given Throable is an ExecutionException, this can return the cause of ExecutionException. */ private RuntimeException throwIfPossible(Throwable t) throws IOException, ChannelException diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java index 4a8ae2f96cd4..b47ec8901693 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java @@ -241,7 +241,10 @@ private boolean possiblyStopAndRemoveSupervisorInternal(String id, boolean write } if (writeTombstone) { - metadataSupervisorManager.insert(id, new NoopSupervisorSpec(null, pair.rhs.getDataSources())); // where NoopSupervisorSpec is a tombstone + metadataSupervisorManager.insert( + id, + new NoopSupervisorSpec(null, pair.rhs.getDataSources()) + ); // where NoopSupervisorSpec is a tombstone } pair.lhs.stop(true); supervisors.remove(id); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamDataSourceMetadata.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamDataSourceMetadata.java new file mode 100644 index 000000000000..2501048c6e73 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamDataSourceMetadata.java @@ -0,0 +1,167 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.indexing.overlord.DataSourceMetadata; +import org.apache.druid.java.util.common.IAE; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +public abstract class SeekableStreamDataSourceMetadata + implements DataSourceMetadata +{ + private final SeekableStreamPartitions seekableStreamPartitions; + + @JsonCreator + public SeekableStreamDataSourceMetadata( + @JsonProperty("partitions") SeekableStreamPartitions seekableStreamPartitions + ) + { + this.seekableStreamPartitions = seekableStreamPartitions; + } + + @JsonProperty("partitions") + public SeekableStreamPartitions getSeekableStreamPartitions() + { + return seekableStreamPartitions; + } + + @Override + public boolean isValidStart() + { + return true; + } + + @Override + public boolean matches(DataSourceMetadata other) + { + if (!getClass().equals(other.getClass())) { + return false; + } + + return plus(other).equals(other.plus(this)); + } + + + @Override + public DataSourceMetadata plus(DataSourceMetadata other) + { + if (!(this.getClass().isInstance(other))) { + throw new IAE( + "Expected instance of %s, got %s", + this.getClass().getCanonicalName(), + other.getClass().getCanonicalName() + ); + } + + @SuppressWarnings("unchecked") + final SeekableStreamDataSourceMetadata that = (SeekableStreamDataSourceMetadata) other; + + if (that.getSeekableStreamPartitions().getStream().equals(seekableStreamPartitions.getStream())) { + // Same stream, merge sequences. + final Map newMap = new HashMap<>(); + + for (Map.Entry entry : seekableStreamPartitions.getPartitionSequenceNumberMap() + .entrySet()) { + newMap.put(entry.getKey(), entry.getValue()); + } + + for (Map.Entry entry : that.getSeekableStreamPartitions() + .getPartitionSequenceNumberMap() + .entrySet()) { + newMap.put(entry.getKey(), entry.getValue()); + } + + return createConcreteDataSourceMetaData(seekableStreamPartitions.getStream(), newMap); + } else { + // Different stream, prefer "other". + return other; + } + } + + + @Override + public DataSourceMetadata minus(DataSourceMetadata other) + { + if (!(this.getClass().isInstance(other))) { + throw new IAE( + "Expected instance of %s, got %s", + this.getClass().getCanonicalName(), + other.getClass().getCanonicalName() + ); + } + + @SuppressWarnings("unchecked") + final SeekableStreamDataSourceMetadata that = (SeekableStreamDataSourceMetadata) other; + + if (that.getSeekableStreamPartitions().getStream().equals(seekableStreamPartitions.getStream())) { + // Same stream, remove partitions present in "that" from "this" + final Map newMap = new HashMap<>(); + + for (Map.Entry entry : seekableStreamPartitions.getPartitionSequenceNumberMap() + .entrySet()) { + if (!that.getSeekableStreamPartitions().getPartitionSequenceNumberMap().containsKey(entry.getKey())) { + newMap.put(entry.getKey(), entry.getValue()); + } + } + + return createConcreteDataSourceMetaData(seekableStreamPartitions.getStream(), newMap); + } else { + // Different stream, prefer "this". + return this; + } + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || !getClass().equals(o.getClass())) { + return false; + } + SeekableStreamDataSourceMetadata that = (SeekableStreamDataSourceMetadata) o; + return Objects.equals(getSeekableStreamPartitions(), that.getSeekableStreamPartitions()); + } + + @Override + public int hashCode() + { + return Objects.hash(getSeekableStreamPartitions()); + } + + @Override + public String toString() + { + return "SeekableStreamDataSourceMetadata{" + + "SeekableStreamPartitions=" + getSeekableStreamPartitions() + + '}'; + } + + protected abstract SeekableStreamDataSourceMetadata createConcreteDataSourceMetaData( + String streamId, + Map newMap + ); +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java new file mode 100644 index 000000000000..7f279917a3f1 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java @@ -0,0 +1,298 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + +import com.fasterxml.jackson.annotation.JacksonInject; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Joiner; +import com.google.common.base.Optional; +import com.google.common.base.Preconditions; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.impl.InputRowParser; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.appenderator.ActionBasedSegmentAllocator; +import org.apache.druid.indexing.appenderator.ActionBasedUsedSegmentChecker; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.actions.SegmentAllocateAction; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.AbstractTask; +import org.apache.druid.indexing.common.task.TaskResource; +import org.apache.druid.indexing.common.task.Tasks; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.utils.RandomIdUtils; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.parsers.ParseException; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.query.NoopQueryRunner; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.realtime.FireDepartmentMetrics; +import org.apache.druid.segment.realtime.appenderator.Appenderator; +import org.apache.druid.segment.realtime.appenderator.Appenderators; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.firehose.ChatHandler; +import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.utils.CircularBuffer; + +import java.nio.ByteBuffer; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ThreadLocalRandom; + + +public abstract class SeekableStreamIndexTask extends AbstractTask + implements ChatHandler +{ + public static final long LOCK_ACQUIRE_TIMEOUT_SECONDS = 15; + private static final Random RANDOM = ThreadLocalRandom.current(); + private static final EmittingLogger log = new EmittingLogger(SeekableStreamIndexTask.class); + + private final SeekableStreamIndexTaskRunner runner; + protected final DataSchema dataSchema; + protected final InputRowParser parser; + protected final SeekableStreamIndexTaskTuningConfig tuningConfig; + protected final SeekableStreamIndexTaskIOConfig ioConfig; + protected final Optional chatHandlerProvider; + protected final Map context; + protected final AuthorizerMapper authorizerMapper; + protected final RowIngestionMetersFactory rowIngestionMetersFactory; + protected final CircularBuffer savedParseExceptions; + + @JsonCreator + public SeekableStreamIndexTask( + @JsonProperty("id") String id, + @JsonProperty("resource") TaskResource taskResource, + @JsonProperty("dataSchema") DataSchema dataSchema, + @JsonProperty("tuningConfig") SeekableStreamIndexTaskTuningConfig tuningConfig, + @JsonProperty("ioConfig") SeekableStreamIndexTaskIOConfig ioConfig, + @JsonProperty("context") Map context, + @JacksonInject ChatHandlerProvider chatHandlerProvider, + @JacksonInject AuthorizerMapper authorizerMapper, + @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory, + String groupId + ) + { + super( + id, + groupId, + taskResource, + dataSchema.getDataSource(), + context + ); + this.dataSchema = Preconditions.checkNotNull(dataSchema, "dataSchema"); + this.parser = Preconditions.checkNotNull((InputRowParser) dataSchema.getParser(), "parser"); + this.tuningConfig = Preconditions.checkNotNull(tuningConfig, "tuningConfig"); + this.ioConfig = Preconditions.checkNotNull(ioConfig, "ioConfig"); + this.chatHandlerProvider = Optional.fromNullable(chatHandlerProvider); + if (tuningConfig.getMaxSavedParseExceptions() > 0) { + savedParseExceptions = new CircularBuffer<>(tuningConfig.getMaxSavedParseExceptions()); + } else { + savedParseExceptions = null; + } + this.context = context; + this.authorizerMapper = authorizerMapper; + this.rowIngestionMetersFactory = rowIngestionMetersFactory; + this.runner = createTaskRunner(); + } + + + private static String makeTaskId(String dataSource, String type) + { + final String suffix = RandomIdUtils.getRandomId(); + return Joiner.on("_").join(type, dataSource, suffix); + } + + protected static String getFormattedId(String dataSource, String type) + { + return makeTaskId(dataSource, type); + } + + protected static String getFormattedGroupId(String dataSource, String type) + { + return StringUtils.format("%s_%s", type, dataSource); + } + + + @Override + public int getPriority() + { + return getContextValue(Tasks.PRIORITY_KEY, Tasks.DEFAULT_REALTIME_TASK_PRIORITY); + } + + @Override + public boolean isReady(TaskActionClient taskActionClient) + { + return true; + } + + @JsonProperty + public DataSchema getDataSchema() + { + return dataSchema; + } + + @JsonProperty + public SeekableStreamIndexTaskTuningConfig getTuningConfig() + { + return tuningConfig; + } + + @JsonProperty("ioConfig") + public SeekableStreamIndexTaskIOConfig getIOConfig() + { + return ioConfig; + } + + @Override + public TaskStatus run(final TaskToolbox toolbox) + { + return runner.run(toolbox); + } + + @Override + public boolean canRestore() + { + return true; + } + + @Override + public void stopGracefully() + { + runner.stopGracefully(); + } + + @Override + public QueryRunner getQueryRunner(Query query) + { + if (runner.getAppenderator() == null) { + // Not yet initialized, no data yet, just return a noop runner. + return new NoopQueryRunner<>(); + } + + return (queryPlus, responseContext) -> queryPlus.run(runner.getAppenderator(), responseContext); + } + + public Appenderator newAppenderator(FireDepartmentMetrics metrics, TaskToolbox toolbox) + { + return Appenderators.createRealtime( + dataSchema, + tuningConfig.withBasePersistDirectory(toolbox.getPersistDir()), + metrics, + toolbox.getSegmentPusher(), + toolbox.getObjectMapper(), + toolbox.getIndexIO(), + toolbox.getIndexMergerV9(), + toolbox.getQueryRunnerFactoryConglomerate(), + toolbox.getSegmentAnnouncer(), + toolbox.getEmitter(), + toolbox.getQueryExecutorService(), + toolbox.getCache(), + toolbox.getCacheConfig(), + toolbox.getCachePopulatorStats() + ); + } + + public StreamAppenderatorDriver newDriver( + final Appenderator appenderator, + final TaskToolbox toolbox, + final FireDepartmentMetrics metrics + ) + { + return new StreamAppenderatorDriver( + appenderator, + new ActionBasedSegmentAllocator( + toolbox.getTaskActionClient(), + dataSchema, + (schema, row, sequenceName, previousSegmentId, skipSegmentLineageCheck) -> new SegmentAllocateAction( + schema.getDataSource(), + row.getTimestamp(), + schema.getGranularitySpec().getQueryGranularity(), + schema.getGranularitySpec().getSegmentGranularity(), + sequenceName, + previousSegmentId, + skipSegmentLineageCheck + ) + ), + toolbox.getSegmentHandoffNotifierFactory(), + new ActionBasedUsedSegmentChecker(toolbox.getTaskActionClient()), + toolbox.getDataSegmentKiller(), + toolbox.getObjectMapper(), + metrics + ); + } + + public boolean withinMinMaxRecordTime(final InputRow row) + { + final boolean beforeMinimumMessageTime = ioConfig.getMinimumMessageTime().isPresent() + && ioConfig.getMinimumMessageTime().get().isAfter(row.getTimestamp()); + + final boolean afterMaximumMessageTime = ioConfig.getMaximumMessageTime().isPresent() + && ioConfig.getMaximumMessageTime().get().isBefore(row.getTimestamp()); + + if (!Intervals.ETERNITY.contains(row.getTimestamp())) { + final String errorMsg = StringUtils.format( + "Encountered row with timestamp that cannot be represented as a long: [%s]", + row + ); + throw new ParseException(errorMsg); + } + + if (log.isDebugEnabled()) { + if (beforeMinimumMessageTime) { + log.debug( + "CurrentTimeStamp[%s] is before MinimumMessageTime[%s]", + row.getTimestamp(), + ioConfig.getMinimumMessageTime().get() + ); + } else if (afterMaximumMessageTime) { + log.debug( + "CurrentTimeStamp[%s] is after MaximumMessageTime[%s]", + row.getTimestamp(), + ioConfig.getMaximumMessageTime().get() + ); + } + } + + return !beforeMinimumMessageTime && !afterMaximumMessageTime; + } + + protected abstract SeekableStreamIndexTaskRunner createTaskRunner(); + + protected abstract RecordSupplier newTaskRecordSupplier(); + + @VisibleForTesting + public Appenderator getAppenderator() + { + return runner.getAppenderator(); + } + + @VisibleForTesting + public SeekableStreamIndexTaskRunner getRunner() + { + return runner; + } + +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClient.java new file mode 100644 index 000000000000..bd946cc6e9c7 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClient.java @@ -0,0 +1,381 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.indexing.common.IndexTaskClient; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.RE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.jackson.JacksonUtils; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.java.util.http.client.response.FullResponseHolder; +import org.jboss.netty.handler.codec.http.HttpMethod; +import org.jboss.netty.handler.codec.http.HttpResponseStatus; +import org.joda.time.DateTime; +import org.joda.time.Duration; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import java.util.TreeMap; + +public abstract class SeekableStreamIndexTaskClient extends IndexTaskClient +{ + private static final EmittingLogger log = new EmittingLogger(SeekableStreamIndexTaskClient.class); + + public SeekableStreamIndexTaskClient( + HttpClient httpClient, + ObjectMapper jsonMapper, + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ) + { + super(httpClient, jsonMapper, taskInfoProvider, httpTimeout, dataSource, numThreads, numRetries); + } + + public boolean stop(final String id, final boolean publish) + { + log.debug("Stop task[%s] publish[%s]", id, publish); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent( + id, + HttpMethod.POST, + "stop", + publish ? "publish=true" : null, + true + ); + return isSuccess(response); + } + catch (NoTaskLocationException e) { + return false; + } + catch (TaskNotRunnableException e) { + log.info("Task [%s] couldn't be stopped because it is no longer running", id); + return true; + } + catch (Exception e) { + log.warn(e, "Exception while stopping task [%s]", id); + return false; + } + } + + public boolean resume(final String id) + { + log.debug("Resume task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.POST, "resume", null, true); + return isSuccess(response); + } + catch (NoTaskLocationException | IOException e) { + log.warn(e, "Exception while stopping task [%s]", id); + return false; + } + } + + public Map pause(final String id) + { + log.debug("Pause task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent( + id, + HttpMethod.POST, + "pause", + null, + true + ); + + final HttpResponseStatus responseStatus = response.getStatus(); + final String responseContent = response.getContent(); + + if (responseStatus.equals(HttpResponseStatus.OK)) { + log.info("Task [%s] paused successfully", id); + return deserializeMap(responseContent, Map.class, getPartitionType(), getSequenceType()); + } else if (responseStatus.equals(HttpResponseStatus.ACCEPTED)) { + // The task received the pause request, but its status hasn't been changed yet. + while (true) { + final SeekableStreamIndexTaskRunner.Status status = getStatus(id); + if (status == SeekableStreamIndexTaskRunner.Status.PAUSED) { + return getCurrentOffsets(id, true); + } + + final Duration delay = newRetryPolicy().getAndIncrementRetryDelay(); + if (delay == null) { + throw new ISE( + "Task [%s] failed to change its status from [%s] to [%s], aborting", + id, + status, + SeekableStreamIndexTaskRunner.Status.PAUSED + ); + } else { + final long sleepTime = delay.getMillis(); + log.info( + "Still waiting for task [%s] to change its status to [%s]; will try again in [%s]", + id, + SeekableStreamIndexTaskRunner.Status.PAUSED, + new Duration(sleepTime).toString() + ); + Thread.sleep(sleepTime); + } + } + } else { + throw new ISE( + "Pause request for task [%s] failed with response [%s] : [%s]", + id, + responseStatus, + responseContent + ); + } + } + catch (NoTaskLocationException e) { + log.error("Exception [%s] while pausing Task [%s]", e.getMessage(), id); + return ImmutableMap.of(); + } + catch (IOException | InterruptedException e) { + throw new RE(e, "Exception [%s] while pausing Task [%s]", e.getMessage(), id); + } + } + + public SeekableStreamIndexTaskRunner.Status getStatus(final String id) + { + log.debug("GetStatus task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "status", null, true); + return deserialize(response.getContent(), SeekableStreamIndexTaskRunner.Status.class); + } + catch (NoTaskLocationException e) { + return SeekableStreamIndexTaskRunner.Status.NOT_STARTED; + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + + @Nullable + public DateTime getStartTime(final String id) + { + log.debug("GetStartTime task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "time/start", null, true); + return response.getContent() == null || response.getContent().isEmpty() + ? null + : deserialize(response.getContent(), DateTime.class); + } + catch (NoTaskLocationException e) { + return null; + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public Map getMovingAverages(final String id) + { + log.debug("GetMovingAverages task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent( + id, + HttpMethod.GET, + "rowStats", + null, + true + ); + return response.getContent() == null || response.getContent().isEmpty() + ? Collections.emptyMap() + : deserialize(response.getContent(), JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT); + } + catch (NoTaskLocationException e) { + return Collections.emptyMap(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public Map getCurrentOffsets(final String id, final boolean retry) + { + log.debug("GetCurrentOffsets task[%s] retry[%s]", id, retry); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent( + id, + HttpMethod.GET, + "offsets/current", + null, + retry + ); + return deserializeMap(response.getContent(), Map.class, getPartitionType(), getSequenceType()); + } + catch (NoTaskLocationException e) { + return ImmutableMap.of(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public TreeMap> getCheckpoints(final String id, final boolean retry) + { + log.debug("GetCheckpoints task[%s] retry[%s]", id, retry); + try { + final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "checkpoints", null, retry); + return deserializeNestedValueMap( + response.getContent(), + TreeMap.class, + Integer.class, + Map.class, + getPartitionType(), + getSequenceType() + ); + } + catch (NoTaskLocationException e) { + return new TreeMap<>(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public ListenableFuture>> getCheckpointsAsync( + final String id, + final boolean retry + ) + { + return doAsync(() -> getCheckpoints(id, retry)); + } + + public Map getEndOffsets(final String id) + { + log.debug("GetEndOffsets task[%s]", id); + + try { + final FullResponseHolder response = submitRequestWithEmptyContent(id, HttpMethod.GET, "offsets/end", null, true); + return deserializeMap(response.getContent(), Map.class, getPartitionType(), getSequenceType()); + } + catch (NoTaskLocationException e) { + return ImmutableMap.of(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + public boolean setEndOffsets( + final String id, + final Map endOffsets, + final boolean finalize + ) throws IOException + { + log.debug("SetEndOffsets task[%s] endOffsets[%s] finalize[%s]", id, endOffsets, finalize); + + try { + final FullResponseHolder response = submitJsonRequest( + id, + HttpMethod.POST, + "offsets/end", + StringUtils.format("finish=%s", finalize), + serialize(endOffsets), + true + ); + return isSuccess(response); + } + catch (NoTaskLocationException e) { + return false; + } + } + + public ListenableFuture stopAsync(final String id, final boolean publish) + { + return doAsync(() -> stop(id, publish)); + } + + + public ListenableFuture resumeAsync(final String id) + { + return doAsync(() -> resume(id)); + } + + + public ListenableFuture getStartTimeAsync(final String id) + { + return doAsync(() -> getStartTime(id)); + } + + + public ListenableFuture> pauseAsync(final String id) + { + return doAsync(() -> pause(id)); + } + + public ListenableFuture setEndOffsetsAsync( + final String id, + final Map endOffsets, + final boolean finalize + ) + { + return doAsync(() -> setEndOffsets(id, endOffsets, finalize)); + } + + public ListenableFuture> getCurrentOffsetsAsync( + final String id, + final boolean retry + ) + { + return doAsync(() -> getCurrentOffsets(id, retry)); + } + + public ListenableFuture> getEndOffsetsAsync(final String id) + { + return doAsync(() -> getEndOffsets(id)); + } + + + public ListenableFuture> getMovingAveragesAsync(final String id) + { + return doAsync(() -> getMovingAverages(id)); + } + + + public ListenableFuture getStatusAsync(final String id) + { + return doAsync(() -> getStatus(id)); + } + + protected abstract Class getPartitionType(); + + protected abstract Class getSequenceType(); +} + + diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java new file mode 100644 index 000000000000..6a909704bdaa --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.Inject; +import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.common.task.IndexTaskClientFactory; +import org.apache.druid.java.util.http.client.HttpClient; +import org.joda.time.Duration; + +public abstract class SeekableStreamIndexTaskClientFactory + implements IndexTaskClientFactory +{ + private HttpClient httpClient; + private ObjectMapper mapper; + + @Inject + public SeekableStreamIndexTaskClientFactory( + @EscalatedGlobal HttpClient httpClient, + @Json ObjectMapper mapper + ) + { + this.httpClient = httpClient; + this.mapper = mapper; + } + + @Override + public abstract T build( + TaskInfoProvider taskInfoProvider, + String dataSource, + int numThreads, + Duration httpTimeout, + long numRetries + ); + + protected HttpClient getHttpClient() + { + return httpClient; + } + + protected ObjectMapper getMapper() + { + return mapper; + } +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIOConfig.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskIOConfig.java similarity index 62% rename from extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIOConfig.java rename to indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskIOConfig.java index 6a9af7fcea94..6c469c7d0123 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIOConfig.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskIOConfig.java @@ -17,7 +17,7 @@ * under the License. */ -package org.apache.druid.indexing.kafka; +package org.apache.druid.indexing.seekablestream; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; @@ -27,9 +27,10 @@ import org.joda.time.DateTime; import javax.annotation.Nullable; -import java.util.Map; +import java.util.Collections; +import java.util.Set; -public class KafkaIOConfig implements IOConfig +public abstract class SeekableStreamIndexTaskIOConfig implements IOConfig { private static final boolean DEFAULT_USE_TRANSACTION = true; private static final boolean DEFAULT_SKIP_OFFSET_GAPS = false; @@ -37,55 +38,60 @@ public class KafkaIOConfig implements IOConfig @Nullable private final Integer taskGroupId; private final String baseSequenceName; - private final KafkaPartitions startPartitions; - private final KafkaPartitions endPartitions; - private final Map consumerProperties; + private final SeekableStreamPartitions startPartitions; + private final SeekableStreamPartitions endPartitions; private final boolean useTransaction; private final Optional minimumMessageTime; private final Optional maximumMessageTime; private final boolean skipOffsetGaps; + private final Set exclusiveStartSequenceNumberPartitions; @JsonCreator - public KafkaIOConfig( + public SeekableStreamIndexTaskIOConfig( @JsonProperty("taskGroupId") @Nullable Integer taskGroupId, // can be null for backward compabitility @JsonProperty("baseSequenceName") String baseSequenceName, - @JsonProperty("startPartitions") KafkaPartitions startPartitions, - @JsonProperty("endPartitions") KafkaPartitions endPartitions, - @JsonProperty("consumerProperties") Map consumerProperties, + @JsonProperty("startPartitions") SeekableStreamPartitions startPartitions, + @JsonProperty("endPartitions") SeekableStreamPartitions endPartitions, @JsonProperty("useTransaction") Boolean useTransaction, @JsonProperty("minimumMessageTime") DateTime minimumMessageTime, @JsonProperty("maximumMessageTime") DateTime maximumMessageTime, - @JsonProperty("skipOffsetGaps") Boolean skipOffsetGaps + @JsonProperty("skipOffsetGaps") Boolean skipOffsetGaps, + @JsonProperty("exclusiveStartSequenceNumberPartitions") + Set exclusiveStartSequenceNumberPartitions ) { this.taskGroupId = taskGroupId; this.baseSequenceName = Preconditions.checkNotNull(baseSequenceName, "baseSequenceName"); this.startPartitions = Preconditions.checkNotNull(startPartitions, "startPartitions"); this.endPartitions = Preconditions.checkNotNull(endPartitions, "endPartitions"); - this.consumerProperties = Preconditions.checkNotNull(consumerProperties, "consumerProperties"); this.useTransaction = useTransaction != null ? useTransaction : DEFAULT_USE_TRANSACTION; this.minimumMessageTime = Optional.fromNullable(minimumMessageTime); this.maximumMessageTime = Optional.fromNullable(maximumMessageTime); this.skipOffsetGaps = skipOffsetGaps != null ? skipOffsetGaps : DEFAULT_SKIP_OFFSET_GAPS; + this.exclusiveStartSequenceNumberPartitions = exclusiveStartSequenceNumberPartitions == null + ? Collections.emptySet() + : exclusiveStartSequenceNumberPartitions; Preconditions.checkArgument( - startPartitions.getTopic().equals(endPartitions.getTopic()), - "start topic and end topic must match" + startPartitions.getStream().equals(endPartitions.getStream()), + "start topic/stream and end topic/stream must match" ); Preconditions.checkArgument( - startPartitions.getPartitionOffsetMap().keySet().equals(endPartitions.getPartitionOffsetMap().keySet()), + startPartitions.getPartitionSequenceNumberMap() + .keySet() + .equals(endPartitions.getPartitionSequenceNumberMap().keySet()), "start partition set and end partition set must match" ); + } - for (int partition : endPartitions.getPartitionOffsetMap().keySet()) { - Preconditions.checkArgument( - endPartitions.getPartitionOffsetMap().get(partition) >= startPartitions.getPartitionOffsetMap() - .get(partition), - "end offset must be >= start offset for partition[%s]", - partition - ); - } + // exclusive starting sequence partitions are used only for kinesis where the starting + // sequence number for certain partitions are discarded because they've already been + // read by a previous task + @JsonProperty + public Set getExclusiveStartSequenceNumberPartitions() + { + return exclusiveStartSequenceNumberPartitions; } @Nullable @@ -102,23 +108,17 @@ public String getBaseSequenceName() } @JsonProperty - public KafkaPartitions getStartPartitions() + public SeekableStreamPartitions getStartPartitions() { return startPartitions; } @JsonProperty - public KafkaPartitions getEndPartitions() + public SeekableStreamPartitions getEndPartitions() { return endPartitions; } - @JsonProperty - public Map getConsumerProperties() - { - return consumerProperties; - } - @JsonProperty public boolean isUseTransaction() { @@ -142,20 +142,4 @@ public boolean isSkipOffsetGaps() { return skipOffsetGaps; } - - @Override - public String toString() - { - return "KafkaIOConfig{" + - "taskGroupId=" + taskGroupId + - ", baseSequenceName='" + baseSequenceName + '\'' + - ", startPartitions=" + startPartitions + - ", endPartitions=" + endPartitions + - ", consumerProperties=" + consumerProperties + - ", useTransaction=" + useTransaction + - ", minimumMessageTime=" + minimumMessageTime + - ", maximumMessageTime=" + maximumMessageTime + - ", skipOffsetGaps=" + skipOffsetGaps + - '}'; - } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java new file mode 100644 index 000000000000..e85de9d8db68 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java @@ -0,0 +1,2042 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Function; +import com.google.common.base.Joiner; +import com.google.common.base.Optional; +import com.google.common.base.Preconditions; +import com.google.common.base.Supplier; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Sets; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.data.input.Committer; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.impl.InputRowParser; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.LookupNodeService; +import org.apache.druid.discovery.NodeType; +import org.apache.druid.indexer.IngestionState; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.IngestionStatsAndErrorsTaskReport; +import org.apache.druid.indexing.common.IngestionStatsAndErrorsTaskReportData; +import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder; +import org.apache.druid.indexing.common.TaskReport; +import org.apache.druid.indexing.common.TaskToolbox; +import org.apache.druid.indexing.common.actions.CheckPointDataSourceMetadataAction; +import org.apache.druid.indexing.common.actions.ResetDataSourceMetadataAction; +import org.apache.druid.indexing.common.actions.SegmentTransactionalInsertAction; +import org.apache.druid.indexing.common.stats.RowIngestionMeters; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.IndexTaskUtils; +import org.apache.druid.indexing.common.task.RealtimeIndexTask; +import org.apache.druid.indexing.seekablestream.common.OrderedPartitionableRecord; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.collect.Utils; +import org.apache.druid.java.util.common.parsers.ParseException; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.segment.indexing.RealtimeIOConfig; +import org.apache.druid.segment.realtime.FireDepartment; +import org.apache.druid.segment.realtime.FireDepartmentMetrics; +import org.apache.druid.segment.realtime.appenderator.Appenderator; +import org.apache.druid.segment.realtime.appenderator.AppenderatorDriverAddResult; +import org.apache.druid.segment.realtime.appenderator.SegmentsAndMetadata; +import org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver; +import org.apache.druid.segment.realtime.appenderator.TransactionalSegmentPublisher; +import org.apache.druid.segment.realtime.firehose.ChatHandler; +import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider; +import org.apache.druid.server.security.Access; +import org.apache.druid.server.security.Action; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.utils.CircularBuffer; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import javax.servlet.http.HttpServletRequest; +import javax.validation.constraints.NotNull; +import javax.ws.rs.Consumes; +import javax.ws.rs.DefaultValue; +import javax.ws.rs.GET; +import javax.ws.rs.POST; +import javax.ws.rs.Path; +import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import java.util.stream.Collectors; + +/** + * Interface for abstracting the indexing task run logic. + * + * @param Partition Number Type + * @param Sequence Number Type + */ +public abstract class SeekableStreamIndexTaskRunner implements ChatHandler +{ + public enum Status + { + NOT_STARTED, + STARTING, + READING, + PAUSED, + PUBLISHING + } + + private static final EmittingLogger log = new EmittingLogger(SeekableStreamIndexTaskRunner.class); + private static final String METADATA_NEXT_PARTITIONS = "nextPartitions"; + private static final String METADATA_PUBLISH_PARTITIONS = "publishPartitions"; + + private final Map endOffsets; + private final Map currOffsets = new ConcurrentHashMap<>(); + private final Map lastPersistedOffsets = new ConcurrentHashMap<>(); + + // The pause lock and associated conditions are to support coordination between the Jetty threads and the main + // ingestion loop. The goal is to provide callers of the API a guarantee that if pause() returns successfully + // the ingestion loop has been stopped at the returned sequences and will not ingest any more data until resumed. The + // fields are used as follows (every step requires acquiring [pauseLock]): + // Pausing: + // - In pause(), [pauseRequested] is set to true and then execution waits for [status] to change to PAUSED, with the + // condition checked when [hasPaused] is signalled. + // - In possiblyPause() called from the main loop, if [pauseRequested] is true, [status] is set to PAUSED, + // [hasPaused] is signalled, and execution pauses until [pauseRequested] becomes false, either by being set or by + // the [pauseMillis] timeout elapsing. [pauseRequested] is checked when [shouldResume] is signalled. + // Resuming: + // - In resume(), [pauseRequested] is set to false, [shouldResume] is signalled, and execution waits for [status] to + // change to something other than PAUSED, with the condition checked when [shouldResume] is signalled. + // - In possiblyPause(), when [shouldResume] is signalled, if [pauseRequested] has become false the pause loop ends, + // [status] is changed to STARTING and [shouldResume] is signalled. + private final Lock pauseLock = new ReentrantLock(); + private final Condition hasPaused = pauseLock.newCondition(); + private final Condition shouldResume = pauseLock.newCondition(); + + protected final AtomicBoolean stopRequested = new AtomicBoolean(false); + private final AtomicBoolean publishOnStop = new AtomicBoolean(false); + + // [statusLock] is used to synchronize the Jetty thread calling stopGracefully() with the main run thread. It prevents + // the main run thread from switching into a publishing state while the stopGracefully() thread thinks it's still in + // a pre-publishing state. This is important because stopGracefully() will try to use the [stopRequested] flag to stop + // the main thread where possible, but this flag is not honored once publishing has begun so in this case we must + // interrupt the thread. The lock ensures that if the run thread is about to transition into publishing state, it + // blocks until after stopGracefully() has set [stopRequested] and then does a final check on [stopRequested] before + // transitioning to publishing state. + private final Object statusLock = new Object(); + + protected final Lock pollRetryLock = new ReentrantLock(); + protected final Condition isAwaitingRetry = pollRetryLock.newCondition(); + + private final SeekableStreamIndexTask task; + private final SeekableStreamIndexTaskIOConfig ioConfig; + private final SeekableStreamIndexTaskTuningConfig tuningConfig; + private final InputRowParser parser; + private final AuthorizerMapper authorizerMapper; + private final Optional chatHandlerProvider; + private final CircularBuffer savedParseExceptions; + private final String stream; + private final RowIngestionMeters rowIngestionMeters; + + private final Set publishingSequences = Sets.newConcurrentHashSet(); + private final List> publishWaitList = new ArrayList<>(); + private final List> handOffWaitList = new ArrayList<>(); + private final Map initialOffsetsSnapshot = new HashMap<>(); + private final Set exclusiveStartingPartitions = new HashSet<>(); + + private volatile DateTime startTime; + private volatile Status status = Status.NOT_STARTED; // this is only ever set by the task runner thread (runThread) + private volatile TaskToolbox toolbox; + private volatile Thread runThread; + private volatile Appenderator appenderator; + private volatile StreamAppenderatorDriver driver; + private volatile IngestionState ingestionState; + + protected volatile boolean pauseRequested = false; + private volatile long nextCheckpointTime; + + private volatile CopyOnWriteArrayList sequences; + private volatile Throwable backgroundThreadException; + + public SeekableStreamIndexTaskRunner( + final SeekableStreamIndexTask task, + final InputRowParser parser, + final AuthorizerMapper authorizerMapper, + final Optional chatHandlerProvider, + final CircularBuffer savedParseExceptions, + final RowIngestionMetersFactory rowIngestionMetersFactory + ) + { + Preconditions.checkNotNull(task); + this.task = task; + this.ioConfig = task.getIOConfig(); + this.tuningConfig = task.getTuningConfig(); + this.parser = parser; + this.authorizerMapper = authorizerMapper; + this.chatHandlerProvider = chatHandlerProvider; + this.savedParseExceptions = savedParseExceptions; + this.stream = ioConfig.getStartPartitions().getStream(); + this.rowIngestionMeters = rowIngestionMetersFactory.createRowIngestionMeters(); + this.endOffsets = new ConcurrentHashMap<>(ioConfig.getEndPartitions().getPartitionSequenceNumberMap()); + this.sequences = new CopyOnWriteArrayList<>(); + this.ingestionState = IngestionState.NOT_STARTED; + + resetNextCheckpointTime(); + } + + + public TaskStatus run(TaskToolbox toolbox) + { + try { + return runInternal(toolbox); + } + catch (Exception e) { + log.error(e, "Encountered exception while running task."); + final String errorMsg = Throwables.getStackTraceAsString(e); + toolbox.getTaskReportFileWriter().write(getTaskCompletionReports(errorMsg)); + return TaskStatus.failure( + task.getId(), + errorMsg + ); + } + } + + private TaskStatus runInternal(TaskToolbox toolbox) throws Exception + { + log.info("SeekableStream indexing task starting up!"); + startTime = DateTimes.nowUtc(); + status = Status.STARTING; + this.toolbox = toolbox; + + + if (!restoreSequences()) { + final TreeMap> checkpoints = getCheckPointsFromContext( + toolbox, + task.getContextValue("checkpoints") + ); + if (checkpoints != null) { + boolean exclusive = false; + Iterator>> sequenceOffsets = checkpoints.entrySet() + .iterator(); + Map.Entry> previous = sequenceOffsets.next(); + while (sequenceOffsets.hasNext()) { + Map.Entry> current = sequenceOffsets.next(); + sequences.add(new SequenceMetadata( + previous.getKey(), + StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), previous.getKey()), + previous.getValue(), + current.getValue(), + true, + exclusive ? previous.getValue().keySet() : null + )); + previous = current; + exclusive = true; + } + sequences.add(new SequenceMetadata( + previous.getKey(), + StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), previous.getKey()), + previous.getValue(), + endOffsets, + false, + exclusive ? previous.getValue().keySet() : null + )); + } else { + sequences.add(new SequenceMetadata( + 0, + StringUtils.format("%s_%s", ioConfig.getBaseSequenceName(), 0), + ioConfig.getStartPartitions().getPartitionSequenceNumberMap(), + endOffsets, + false, + null + )); + } + } + + log.info("Starting with sequences: %s", sequences); + + if (chatHandlerProvider.isPresent()) { + log.info("Found chat handler of class[%s]", chatHandlerProvider.get().getClass().getName()); + chatHandlerProvider.get().register(task.getId(), this, false); + } else { + log.warn("No chat handler detected"); + } + + runThread = Thread.currentThread(); + + // Set up FireDepartmentMetrics + final FireDepartment fireDepartmentForMetrics = new FireDepartment( + task.getDataSchema(), + new RealtimeIOConfig(null, null, null), + null + ); + FireDepartmentMetrics fireDepartmentMetrics = fireDepartmentForMetrics.getMetrics(); + toolbox.getMonitorScheduler() + .addMonitor(TaskRealtimeMetricsMonitorBuilder.build(task, fireDepartmentForMetrics, rowIngestionMeters)); + + final String lookupTier = task.getContextValue(RealtimeIndexTask.CTX_KEY_LOOKUP_TIER); + final LookupNodeService lookupNodeService = lookupTier == null ? + toolbox.getLookupNodeService() : + new LookupNodeService(lookupTier); + + final DiscoveryDruidNode discoveryDruidNode = new DiscoveryDruidNode( + toolbox.getDruidNode(), + NodeType.PEON, + ImmutableMap.of( + toolbox.getDataNodeService().getName(), toolbox.getDataNodeService(), + lookupNodeService.getName(), lookupNodeService + ) + ); + + Throwable caughtExceptionOuter = null; + try (final RecordSupplier recordSupplier = task.newTaskRecordSupplier()) { + toolbox.getDataSegmentServerAnnouncer().announce(); + toolbox.getDruidNodeAnnouncer().announce(discoveryDruidNode); + + appenderator = task.newAppenderator(fireDepartmentMetrics, toolbox); + driver = task.newDriver(appenderator, toolbox, fireDepartmentMetrics); + + final String stream = ioConfig.getStartPartitions().getStream(); + + // Start up, set up initial sequences. + final Object restoredMetadata = driver.startJob(); + if (restoredMetadata == null) { + // no persist has happened so far + // so either this is a brand new task or replacement of a failed task + Preconditions.checkState(sequences.get(0).startOffsets.entrySet().stream().allMatch( + partitionOffsetEntry -> + createSequenceNumber(partitionOffsetEntry.getValue()).compareTo( + createSequenceNumber(ioConfig.getStartPartitions() + .getPartitionSequenceNumberMap() + .get(partitionOffsetEntry.getKey()) + )) >= 0 + ), "Sequence sequences are not compatible with start sequences of task"); + currOffsets.putAll(sequences.get(0).startOffsets); + } else { + @SuppressWarnings("unchecked") + final Map restoredMetadataMap = (Map) restoredMetadata; + final SeekableStreamPartitions restoredNextPartitions = deserializeSeekableStreamPartitionsFromMetadata( + toolbox.getObjectMapper(), + restoredMetadataMap.get(METADATA_NEXT_PARTITIONS) + ); + + currOffsets.putAll(restoredNextPartitions.getPartitionSequenceNumberMap()); + + // Sanity checks. + if (!restoredNextPartitions.getStream().equals(ioConfig.getStartPartitions().getStream())) { + throw new ISE( + "WTF?! Restored stream[%s] but expected stream[%s]", + restoredNextPartitions.getStream(), + ioConfig.getStartPartitions().getStream() + ); + } + + if (!currOffsets.keySet().equals(ioConfig.getStartPartitions().getPartitionSequenceNumberMap().keySet())) { + throw new ISE( + "WTF?! Restored partitions[%s] but expected partitions[%s]", + currOffsets.keySet(), + ioConfig.getStartPartitions().getPartitionSequenceNumberMap().keySet() + ); + } + // sequences size can be 0 only when all sequences got published and task stopped before it could finish + // which is super rare + if (sequences.size() == 0 || sequences.get(sequences.size() - 1).isCheckpointed()) { + this.endOffsets.putAll(sequences.size() == 0 + ? currOffsets + : sequences.get(sequences.size() - 1).getEndOffsets()); + log.info("End sequences changed to [%s]", endOffsets); + } + } + + // Filter out partitions with END_OF_SHARD markers since these partitions have already been fully read. This + // should have been done by the supervisor already so this is defensive. + int numPreFilterPartitions = currOffsets.size(); + if (currOffsets.entrySet().removeIf(x -> isEndOfShard(x.getValue()))) { + log.info( + "Removed [%d] partitions from assignment which have already been closed", + numPreFilterPartitions - currOffsets.size() + ); + } + + // Set up committer. + final Supplier committerSupplier = () -> { + final Map snapshot = ImmutableMap.copyOf(currOffsets); + lastPersistedOffsets.clear(); + lastPersistedOffsets.putAll(snapshot); + + return new Committer() + { + @Override + public Object getMetadata() + { + return ImmutableMap.of( + METADATA_NEXT_PARTITIONS, new SeekableStreamPartitions<>( + ioConfig.getStartPartitions().getStream(), + snapshot + ) + ); + } + + @Override + public void run() + { + // Do nothing. + } + }; + }; + + // restart publishing of sequences (if any) + maybePersistAndPublishSequences(committerSupplier); + + Set> assignment = assignPartitions(recordSupplier); + possiblyResetDataSourceMetadata(toolbox, recordSupplier, assignment, currOffsets); + seekToStartingSequence(recordSupplier, assignment); + + ingestionState = IngestionState.BUILD_SEGMENTS; + + // Main loop. + // Could eventually support leader/follower mode (for keeping replicas more in sync) + boolean stillReading = !assignment.isEmpty(); + status = Status.READING; + Throwable caughtExceptionInner = null; + + initialOffsetsSnapshot.putAll(currOffsets); + exclusiveStartingPartitions.addAll(ioConfig.getExclusiveStartSequenceNumberPartitions()); + + try { + while (stillReading) { + if (possiblyPause()) { + // The partition assignments may have changed while paused by a call to setEndOffsets() so reassign + // partitions upon resuming. This is safe even if the end sequences have not been modified. + assignment = assignPartitions(recordSupplier); + possiblyResetDataSourceMetadata(toolbox, recordSupplier, assignment, currOffsets); + seekToStartingSequence(recordSupplier, assignment); + + if (assignment.isEmpty()) { + log.info("All partitions have been fully read"); + publishOnStop.set(true); + stopRequested.set(true); + } + } + + // if stop is requested or task's end sequence is set by call to setEndOffsets method with finish set to true + if (stopRequested.get() || sequences.get(sequences.size() - 1).isCheckpointed()) { + status = Status.PUBLISHING; + } + + if (stopRequested.get()) { + break; + } + + if (backgroundThreadException != null) { + throw new RuntimeException(backgroundThreadException); + } + + checkPublishAndHandoffFailure(); + + maybePersistAndPublishSequences(committerSupplier); + + + // calling getRecord() ensures that exceptions specific to kafka/kinesis like OffsetOutOfRangeException + // are handled in the subclasses. + List> records = getRecords( + recordSupplier, + toolbox + ); + + // note: getRecords() also updates assignment + stillReading = !assignment.isEmpty(); + + SequenceMetadata sequenceToCheckpoint = null; + for (OrderedPartitionableRecord record : records) { + + // for Kafka, the end offsets are exclusive, so skip it + if (isEndSequenceOffsetsExclusive() && + createSequenceNumber(record.getSequenceNumber()).compareTo( + createSequenceNumber(endOffsets.get(record.getPartitionId()))) == 0) { + continue; + } + + // for the first message we receive, check that we were given a message with a sequenceNumber that matches our + // expected starting sequenceNumber + if (!verifyInitialRecordAndSkipExclusivePartition(record, initialOffsetsSnapshot)) { + continue; + } + + log.trace( + "Got stream[%s] partition[%s] sequence[%s].", + record.getStream(), + record.getPartitionId(), + record.getSequenceNumber() + ); + + if (isEndOfShard(record.getSequenceNumber())) { + // shard is closed, applies to Kinesis only + currOffsets.put(record.getPartitionId(), record.getSequenceNumber()); + } else if (createSequenceNumber(record.getSequenceNumber()).compareTo( + createSequenceNumber(endOffsets.get(record.getPartitionId()))) <= 0) { + + + if (!record.getSequenceNumber().equals(currOffsets.get(record.getPartitionId())) + && !ioConfig.isSkipOffsetGaps()) { + throw new ISE( + "WTF?! Got sequence[%s] after sequence[%s] in partition[%s].", + record.getSequenceNumber(), + currOffsets.get(record.getPartitionId()), + record.getPartitionId() + ); + } + + try { + final List valueBytess = record.getData(); + final List rows; + if (valueBytess == null || valueBytess.isEmpty()) { + rows = Utils.nullableListOf((InputRow) null); + } else { + rows = new ArrayList<>(); + for (byte[] valueBytes : valueBytess) { + rows.addAll(parser.parseBatch(ByteBuffer.wrap(valueBytes))); + } + } + boolean isPersistRequired = false; + + final SequenceMetadata sequenceToUse = sequences + .stream() + .filter(sequenceMetadata -> sequenceMetadata.canHandle(record)) + .findFirst() + .orElse(null); + + if (sequenceToUse == null) { + throw new ISE( + "WTH?! cannot find any valid sequence for record with partition [%d] and sequence [%d]. Current sequences: %s", + record.getPartitionId(), + record.getSequenceNumber(), + sequences + ); + } + + for (InputRow row : rows) { + if (row != null && task.withinMinMaxRecordTime(row)) { + final AppenderatorDriverAddResult addResult = driver.add( + row, + sequenceToUse.getSequenceName(), + committerSupplier, + true, + // do not allow incremental persists to happen until all the rows from this batch + // of rows are indexed + false + ); + + if (addResult.isOk()) { + // If the number of rows in the segment exceeds the threshold after adding a row, + // move the segment out from the active segments of BaseAppenderatorDriver to make a new segment. + if (addResult.isPushRequired(tuningConfig) && !sequenceToUse.isCheckpointed()) { + sequenceToCheckpoint = sequenceToUse; + } + isPersistRequired |= addResult.isPersistRequired(); + } else { + // Failure to allocate segment puts determinism at risk, bail out to be safe. + // May want configurable behavior here at some point. + // If we allow continuing, then consider blacklisting the interval for a while to avoid constant checks. + throw new ISE("Could not allocate segment for row with timestamp[%s]", row.getTimestamp()); + } + + if (addResult.getParseException() != null) { + handleParseException(addResult.getParseException(), record); + } else { + rowIngestionMeters.incrementProcessed(); + } + } else { + rowIngestionMeters.incrementThrownAway(); + } + } + if (isPersistRequired) { + Futures.addCallback( + driver.persistAsync(committerSupplier.get()), + new FutureCallback() + { + @Override + public void onSuccess(@Nullable Object result) + { + log.info("Persist completed with metadata [%s]", result); + } + + @Override + public void onFailure(Throwable t) + { + log.error("Persist failed, dying"); + backgroundThreadException = t; + } + } + ); + } + } + catch (ParseException e) { + handleParseException(e, record); + } + + // in kafka, we can easily get the next offset by adding 1, but for kinesis, there's no way + // to get the next sequence number without having to make an expensive api call. So the behavior + // here for kafka is to +1 while for kinesis we simply save the current sequence number + currOffsets.put(record.getPartitionId(), getSequenceNumberToStoreAfterRead(record.getSequenceNumber())); + } + + if ((currOffsets.get(record.getPartitionId()).equals(endOffsets.get(record.getPartitionId())) + || isEndOfShard(currOffsets.get(record.getPartitionId()))) + && assignment.remove(record.getStreamPartition())) { + log.info("Finished reading stream[%s], partition[%s].", record.getStream(), record.getPartitionId()); + recordSupplier.assign(assignment); + stillReading = !assignment.isEmpty(); + } + } + + if (System.currentTimeMillis() > nextCheckpointTime) { + sequenceToCheckpoint = sequences.get(sequences.size() - 1); + } + + if (sequenceToCheckpoint != null && stillReading) { + Preconditions.checkArgument( + sequences.get(sequences.size() - 1) + .getSequenceName() + .equals(sequenceToCheckpoint.getSequenceName()), + "Cannot checkpoint a sequence [%s] which is not the latest one, sequences %s", + sequenceToCheckpoint, + sequences + ); + requestPause(); + final CheckPointDataSourceMetadataAction checkpointAction = new CheckPointDataSourceMetadataAction( + task.getDataSource(), + ioConfig.getTaskGroupId(), + task.getIOConfig().getBaseSequenceName(), + createDataSourceMetadata(new SeekableStreamPartitions<>( + stream, + sequenceToCheckpoint.getStartOffsets() + )), + createDataSourceMetadata(new SeekableStreamPartitions<>(stream, currOffsets)) + ); + if (!toolbox.getTaskActionClient().submit(checkpointAction)) { + throw new ISE("Checkpoint request with sequences [%s] failed, dying", currOffsets); + } + } + } + ingestionState = IngestionState.COMPLETED; + } + catch (Exception e) { + // (1) catch all exceptions while reading from kafka + caughtExceptionInner = e; + log.error(e, "Encountered exception in run() before persisting."); + throw e; + } + finally { + log.info("Persisting all pending data"); + try { + driver.persist(committerSupplier.get()); // persist pending data + } + catch (Exception e) { + if (caughtExceptionInner != null) { + caughtExceptionInner.addSuppressed(e); + } else { + throw e; + } + } + } + + synchronized (statusLock) { + if (stopRequested.get() && !publishOnStop.get()) { + throw new InterruptedException("Stopping without publishing"); + } + + status = Status.PUBLISHING; + } + + for (SequenceMetadata sequenceMetadata : sequences) { + if (!publishingSequences.contains(sequenceMetadata.getSequenceName())) { + // this is done to prevent checks in sequence specific commit supplier from failing + sequenceMetadata.setEndOffsets(currOffsets); + sequenceMetadata.updateAssignments(currOffsets); + publishingSequences.add(sequenceMetadata.getSequenceName()); + // persist already done in finally, so directly add to publishQueue + publishAndRegisterHandoff(sequenceMetadata); + } + } + + if (backgroundThreadException != null) { + throw new RuntimeException(backgroundThreadException); + } + + // Wait for publish futures to complete. + Futures.allAsList(publishWaitList).get(); + + // Wait for handoff futures to complete. + // Note that every publishing task (created by calling AppenderatorDriver.publish()) has a corresponding + // handoffFuture. handoffFuture can throw an exception if 1) the corresponding publishFuture failed or 2) it + // failed to persist sequences. It might also return null if handoff failed, but was recoverable. + // See publishAndRegisterHandoff() for details. + List handedOffList = Collections.emptyList(); + if (tuningConfig.getHandoffConditionTimeout() == 0) { + handedOffList = Futures.allAsList(handOffWaitList).get(); + } else { + try { + handedOffList = Futures.allAsList(handOffWaitList) + .get(tuningConfig.getHandoffConditionTimeout(), TimeUnit.MILLISECONDS); + } + catch (TimeoutException e) { + // Handoff timeout is not an indexing failure, but coordination failure. We simply ignore timeout exception + // here. + log.makeAlert("Timed out after [%d] millis waiting for handoffs", tuningConfig.getHandoffConditionTimeout()) + .addData("TaskId", task.getId()) + .emit(); + } + } + + for (SegmentsAndMetadata handedOff : handedOffList) { + log.info( + "Handoff completed for segments[%s] with metadata[%s].", + Joiner.on(", ").join( + handedOff.getSegments().stream().map(DataSegment::getIdentifier).collect(Collectors.toList()) + ), + Preconditions.checkNotNull(handedOff.getCommitMetadata(), "commitMetadata") + ); + } + + appenderator.close(); + } + catch (InterruptedException | RejectedExecutionException e) { + // (2) catch InterruptedException and RejectedExecutionException thrown for the whole ingestion steps including + // the final publishing. + caughtExceptionOuter = e; + try { + Futures.allAsList(publishWaitList).cancel(true); + Futures.allAsList(handOffWaitList).cancel(true); + if (appenderator != null) { + appenderator.closeNow(); + } + } + catch (Exception e2) { + e.addSuppressed(e2); + } + + // handle the InterruptedException that gets wrapped in a RejectedExecutionException + if (e instanceof RejectedExecutionException + && (e.getCause() == null || !(e.getCause() instanceof InterruptedException))) { + throw e; + } + + // if we were interrupted because we were asked to stop, handle the exception and return success, else rethrow + if (!stopRequested.get()) { + Thread.currentThread().interrupt(); + throw e; + } + + log.info("The task was asked to stop before completing"); + } + catch (Exception e) { + // (3) catch all other exceptions thrown for the whole ingestion steps including the final publishing. + caughtExceptionOuter = e; + try { + Futures.allAsList(publishWaitList).cancel(true); + Futures.allAsList(handOffWaitList).cancel(true); + if (appenderator != null) { + appenderator.closeNow(); + } + } + catch (Exception e2) { + e.addSuppressed(e2); + } + throw e; + } + finally { + try { + + if (driver != null) { + driver.close(); + } + if (chatHandlerProvider.isPresent()) { + chatHandlerProvider.get().unregister(task.getId()); + } + + toolbox.getDruidNodeAnnouncer().unannounce(discoveryDruidNode); + toolbox.getDataSegmentServerAnnouncer().unannounce(); + } + catch (Exception e) { + if (caughtExceptionOuter != null) { + caughtExceptionOuter.addSuppressed(e); + } else { + throw e; + } + } + } + + toolbox.getTaskReportFileWriter().write(getTaskCompletionReports(null)); + return TaskStatus.success(task.getId()); + } + + /** + * checks if the input seqNum marks end of shard. Used by Kinesis only + */ + protected abstract boolean isEndOfShard(SequenceOffsetType seqNum); + + private void checkPublishAndHandoffFailure() throws ExecutionException, InterruptedException + { + // Check if any publishFuture failed. + final List> publishFinished = publishWaitList + .stream() + .filter(Future::isDone) + .collect(Collectors.toList()); + + for (ListenableFuture publishFuture : publishFinished) { + // If publishFuture failed, the below line will throw an exception and catched by (1), and then (2) or (3). + publishFuture.get(); + } + + publishWaitList.removeAll(publishFinished); + + // Check if any handoffFuture failed. + final List> handoffFinished = handOffWaitList + .stream() + .filter(Future::isDone) + .collect(Collectors.toList()); + + for (ListenableFuture handoffFuture : handoffFinished) { + // If handoffFuture failed, the below line will throw an exception and catched by (1), and then (2) or (3). + handoffFuture.get(); + } + + handOffWaitList.removeAll(handoffFinished); + } + + private void publishAndRegisterHandoff(SequenceMetadata sequenceMetadata) + { + log.info("Publishing segments for sequence [%s]", sequenceMetadata); + + final ListenableFuture publishFuture = Futures.transform( + driver.publish( + sequenceMetadata.createPublisher(toolbox, ioConfig.isUseTransaction()), + sequenceMetadata.getCommitterSupplier(stream, lastPersistedOffsets).get(), + Collections.singletonList(sequenceMetadata.getSequenceName()) + ), + (Function) publishedSegmentsAndMetadata -> { + if (publishedSegmentsAndMetadata == null) { + throw new ISE( + "Transaction failure publishing segments for sequence [%s]", + sequenceMetadata + ); + } else { + return publishedSegmentsAndMetadata; + } + } + ); + publishWaitList.add(publishFuture); + + // Create a handoffFuture for every publishFuture. The created handoffFuture must fail if publishFuture fails. + final SettableFuture handoffFuture = SettableFuture.create(); + handOffWaitList.add(handoffFuture); + + Futures.addCallback( + publishFuture, + new FutureCallback() + { + @Override + public void onSuccess(SegmentsAndMetadata publishedSegmentsAndMetadata) + { + log.info( + "Published segments[%s] with metadata[%s].", + publishedSegmentsAndMetadata.getSegments() + .stream() + .map(DataSegment::getIdentifier) + .collect(Collectors.toList()), + Preconditions.checkNotNull(publishedSegmentsAndMetadata.getCommitMetadata(), "commitMetadata") + ); + + sequences.remove(sequenceMetadata); + publishingSequences.remove(sequenceMetadata.getSequenceName()); + try { + persistSequences(); + } + catch (IOException e) { + log.error(e, "Unable to persist state, dying"); + handoffFuture.setException(e); + throw new RuntimeException(e); + } + + Futures.transform( + driver.registerHandoff(publishedSegmentsAndMetadata), + new Function() + { + @Nullable + @Override + public Void apply(@Nullable SegmentsAndMetadata handoffSegmentsAndMetadata) + { + if (handoffSegmentsAndMetadata == null) { + log.warn( + "Failed to handoff segments[%s]", + publishedSegmentsAndMetadata.getSegments() + .stream() + .map(DataSegment::getIdentifier) + .collect(Collectors.toList()) + ); + } + handoffFuture.set(handoffSegmentsAndMetadata); + return null; + } + } + ); + } + + @Override + public void onFailure(Throwable t) + { + log.error(t, "Error while publishing segments for sequence[%s]", sequenceMetadata); + handoffFuture.setException(t); + } + } + ); + } + + private static File getSequencesPersistFile(TaskToolbox toolbox) + { + return new File(toolbox.getPersistDir(), "sequences.json"); + } + + private boolean restoreSequences() throws IOException + { + final File sequencesPersistFile = getSequencesPersistFile(toolbox); + if (sequencesPersistFile.exists()) { + sequences = new CopyOnWriteArrayList<>( + toolbox.getObjectMapper().>readValue( + sequencesPersistFile, + new TypeReference>() + { + } + ) + ); + return true; + } else { + return false; + } + } + + private synchronized void persistSequences() throws IOException + { + log.info("Persisting Sequences Metadata [%s]", sequences); + toolbox.getObjectMapper().writerWithType( + new TypeReference>() + { + } + ).writeValue(getSequencesPersistFile(toolbox), sequences); + } + + private Map getTaskCompletionReports(@Nullable String errorMsg) + { + return TaskReport.buildTaskReports( + new IngestionStatsAndErrorsTaskReport( + task.getId(), + new IngestionStatsAndErrorsTaskReportData( + ingestionState, + getTaskCompletionUnparseableEvents(), + getTaskCompletionRowStats(), + errorMsg + ) + ) + ); + } + + private Map getTaskCompletionUnparseableEvents() + { + Map unparseableEventsMap = new HashMap<>(); + List buildSegmentsParseExceptionMessages = IndexTaskUtils.getMessagesFromSavedParseExceptions( + savedParseExceptions + ); + if (buildSegmentsParseExceptionMessages != null) { + unparseableEventsMap.put(RowIngestionMeters.BUILD_SEGMENTS, buildSegmentsParseExceptionMessages); + } + return unparseableEventsMap; + } + + private Map getTaskCompletionRowStats() + { + Map metrics = new HashMap<>(); + metrics.put( + RowIngestionMeters.BUILD_SEGMENTS, + rowIngestionMeters.getTotals() + ); + return metrics; + } + + + private void maybePersistAndPublishSequences(Supplier committerSupplier) + throws InterruptedException + { + for (SequenceMetadata sequenceMetadata : sequences) { + sequenceMetadata.updateAssignments(currOffsets); + if (!sequenceMetadata.isOpen() && !publishingSequences.contains(sequenceMetadata.getSequenceName())) { + publishingSequences.add(sequenceMetadata.getSequenceName()); + try { + Object result = driver.persist(committerSupplier.get()); + log.info( + "Persist completed with results: [%s], adding sequence [%s] to publish queue", + result, + sequenceMetadata + ); + publishAndRegisterHandoff(sequenceMetadata); + } + catch (InterruptedException e) { + log.warn("Interrupted while persisting sequence [%s]", sequenceMetadata); + throw e; + } + } + } + } + + private Set> assignPartitions( + RecordSupplier recordSupplier + ) + { + final Set> assignment = new HashSet<>(); + for (Map.Entry entry : currOffsets.entrySet()) { + final SequenceOffsetType endOffset = endOffsets.get(entry.getKey()); + if (isEndOfShard(endOffset) + || SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER.equals(endOffset) + || createSequenceNumber(entry.getValue()).compareTo(createSequenceNumber(endOffset)) < 0) { + assignment.add(StreamPartition.of(stream, entry.getKey())); + } else if (entry.getValue().equals(endOffset)) { + log.info("Finished reading partition[%s].", entry.getKey()); + } else { + throw new ISE( + "WTF?! Cannot start from sequence[%,d] > endOffset[%,d]", + entry.getValue(), + endOffset + ); + } + } + + recordSupplier.assign(assignment); + + return assignment; + } + + + private void seekToStartingSequence( + RecordSupplier recordSupplier, + Set> partitions + ) throws InterruptedException + { + for (final StreamPartition partition : partitions) { + final SequenceOffsetType sequence = currOffsets.get(partition.getPartitionId()); + log.info("Seeking partition[%s] to sequence[%s].", partition.getPartitionId(), sequence); + recordSupplier.seek(partition, sequence); + } + } + + /** + * Checks if the pauseRequested flag was set and if so blocks: + * a) if pauseMillis == PAUSE_FOREVER, until pauseRequested is cleared + * b) if pauseMillis != PAUSE_FOREVER, until pauseMillis elapses -or- pauseRequested is cleared + *

+ * If pauseMillis is changed while paused, the new pause timeout will be applied. This allows adjustment of the + * pause timeout (making a timed pause into an indefinite pause and vice versa is valid) without having to resume + * and ensures that the loop continues to stay paused without ingesting any new events. You will need to signal + * shouldResume after adjusting pauseMillis for the new value to take effect. + *

+ * Sets paused = true and signals paused so callers can be notified when the pause command has been accepted. + *

+ * Additionally, pauses if all partitions assignments have been read and pauseAfterRead flag is set. + * + * @return true if a pause request was handled, false otherwise + */ + private boolean possiblyPause() throws InterruptedException + { + pauseLock.lockInterruptibly(); + try { + if (pauseRequested) { + status = Status.PAUSED; + hasPaused.signalAll(); + + while (pauseRequested) { + log.info("Pausing ingestion until resumed"); + shouldResume.await(); + } + + status = Status.READING; + shouldResume.signalAll(); + log.info("Ingestion loop resumed"); + return true; + } + } + finally { + pauseLock.unlock(); + } + + return false; + } + + + private void handleParseException(ParseException pe, OrderedPartitionableRecord record) + { + if (pe.isFromPartiallyValidRow()) { + rowIngestionMeters.incrementProcessedWithError(); + } else { + rowIngestionMeters.incrementUnparseable(); + } + + if (tuningConfig.isLogParseExceptions()) { + log.error( + pe, + "Encountered parse exception on row from partition[%s] sequence[%s]", + record.getPartitionId(), + record.getSequenceNumber() + ); + } + + if (savedParseExceptions != null) { + savedParseExceptions.add(pe); + } + + if (rowIngestionMeters.getUnparseable() + rowIngestionMeters.getProcessedWithError() + > tuningConfig.getMaxParseExceptions()) { + log.error("Max parse exceptions exceeded, terminating task..."); + throw new RuntimeException("Max parse exceptions exceeded, terminating task..."); + } + } + + private boolean isPaused() + { + return status == Status.PAUSED; + } + + private void requestPause() + { + pauseRequested = true; + } + + + protected void sendResetRequestAndWait( + Map, SequenceOffsetType> outOfRangePartitions, + TaskToolbox taskToolbox + ) + throws IOException + { + Map partitionOffsetMap = outOfRangePartitions + .entrySet().stream().collect(Collectors.toMap(x -> x.getKey().getPartitionId(), Map.Entry::getValue)); + + boolean result = taskToolbox + .getTaskActionClient() + .submit( + new ResetDataSourceMetadataAction( + task.getDataSource(), + createDataSourceMetadata( + new SeekableStreamPartitions<>( + ioConfig.getStartPartitions().getStream(), + partitionOffsetMap + ) + ) + ) + ); + + if (result) { + log.makeAlert("Resetting sequences for datasource [%s]", task.getDataSource()) + .addData("partitions", partitionOffsetMap.keySet()) + .emit(); + + requestPause(); + } else { + log.makeAlert("Failed to send reset request for partitions [%s]", partitionOffsetMap.keySet()).emit(); + } + } + + /** + * Authorizes action to be performed on this task's datasource + * + * @return authorization result + */ + private Access authorizationCheck(final HttpServletRequest req, Action action) + { + return IndexTaskUtils.datasourceAuthorizationCheck(req, action, task.getDataSource(), authorizerMapper); + } + + public Appenderator getAppenderator() + { + return appenderator; + } + + @VisibleForTesting + public RowIngestionMeters getRowIngestionMeters() + { + return rowIngestionMeters; + } + + + public void stopGracefully() + { + log.info("Stopping gracefully (status: [%s])", status); + stopRequested.set(true); + + synchronized (statusLock) { + if (status == Status.PUBLISHING) { + runThread.interrupt(); + return; + } + } + + try { + if (pauseLock.tryLock(SeekableStreamIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + try { + if (pauseRequested) { + pauseRequested = false; + shouldResume.signalAll(); + } + } + finally { + pauseLock.unlock(); + } + } else { + log.warn("While stopping: failed to acquire pauseLock before timeout, interrupting run thread"); + runThread.interrupt(); + return; + } + + if (pollRetryLock.tryLock(SeekableStreamIndexTask.LOCK_ACQUIRE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + try { + isAwaitingRetry.signalAll(); + } + finally { + pollRetryLock.unlock(); + } + } else { + log.warn("While stopping: failed to acquire pollRetryLock before timeout, interrupting run thread"); + runThread.interrupt(); + } + } + catch (Exception e) { + Throwables.propagate(e); + } + } + + @POST + @Path("/stop") + public Response stop(@Context final HttpServletRequest req) + { + authorizationCheck(req, Action.WRITE); + stopGracefully(); + return Response.status(Response.Status.OK).build(); + } + + @GET + @Path("/status") + @Produces(MediaType.APPLICATION_JSON) + public Status getStatusHTTP(@Context final HttpServletRequest req) + { + authorizationCheck(req, Action.READ); + return status; + } + + @VisibleForTesting + public Status getStatus() + { + return status; + } + + @GET + @Path("/offsets/current") + @Produces(MediaType.APPLICATION_JSON) + public Map getCurrentOffsets(@Context final HttpServletRequest req) + { + authorizationCheck(req, Action.READ); + return getCurrentOffsets(); + } + + public Map getCurrentOffsets() + { + return currOffsets; + } + + @GET + @Path("/offsets/end") + @Produces(MediaType.APPLICATION_JSON) + public Map getEndOffsetsHTTP(@Context final HttpServletRequest req) + { + authorizationCheck(req, Action.READ); + return getEndOffsets(); + } + + public Map getEndOffsets() + { + return endOffsets; + } + + @POST + @Path("/offsets/end") + @Consumes(MediaType.APPLICATION_JSON) + @Produces(MediaType.APPLICATION_JSON) + public Response setEndOffsetsHTTP( + Map sequences, + @QueryParam("finish") @DefaultValue("true") final boolean finish, + // this field is only for internal purposes, shouldn't be usually set by users + @Context final HttpServletRequest req + ) throws InterruptedException + { + authorizationCheck(req, Action.WRITE); + return setEndOffsets(sequences, finish); + } + + @GET + @Path("/rowStats") + @Produces(MediaType.APPLICATION_JSON) + public Response getRowStats( + @Context final HttpServletRequest req + ) + { + authorizationCheck(req, Action.READ); + Map returnMap = new HashMap<>(); + Map totalsMap = new HashMap<>(); + Map averagesMap = new HashMap<>(); + + totalsMap.put( + RowIngestionMeters.BUILD_SEGMENTS, + rowIngestionMeters.getTotals() + ); + averagesMap.put( + RowIngestionMeters.BUILD_SEGMENTS, + rowIngestionMeters.getMovingAverages() + ); + + returnMap.put("movingAverages", averagesMap); + returnMap.put("totals", totalsMap); + return Response.ok(returnMap).build(); + } + + @GET + @Path("/unparseableEvents") + @Produces(MediaType.APPLICATION_JSON) + public Response getUnparseableEvents( + @Context final HttpServletRequest req + ) + { + authorizationCheck(req, Action.READ); + List events = IndexTaskUtils.getMessagesFromSavedParseExceptions(savedParseExceptions); + return Response.ok(events).build(); + } + + @VisibleForTesting + public Response setEndOffsets( + Map sequenceNumbers, + boolean finish // this field is only for internal purposes, shouldn't be usually set by users + ) throws InterruptedException + { + if (sequenceNumbers == null) { + return Response.status(Response.Status.BAD_REQUEST) + .entity("Request body must contain a map of { partition:endOffset }") + .build(); + } else if (!endOffsets.keySet().containsAll(sequenceNumbers.keySet())) { + return Response.status(Response.Status.BAD_REQUEST) + .entity( + StringUtils.format( + "Request contains partitions not being handled by this task, my partitions: %s", + endOffsets.keySet() + ) + ) + .build(); + } else { + try { + pauseLock.lockInterruptibly(); + // Perform all sequence related checks before checking for isPaused() + // and after acquiring pauseLock to correctly guard against duplicate requests + Preconditions.checkState(sequenceNumbers.size() > 0, "WTH?! No Sequences found to set end sequences"); + + final SequenceMetadata latestSequence = sequences.get(sequences.size() - 1); + // if a partition has not been read yet (contained in initialOffsetsSnapshot), then + // do not mark the starting sequence number as exclusive + Set exclusivePartitions = sequenceNumbers.keySet() + .stream() + .filter(x -> !initialOffsetsSnapshot.containsKey(x) + || ioConfig.getExclusiveStartSequenceNumberPartitions() + .contains(x)) + .collect(Collectors.toSet()); + + if ((latestSequence.getStartOffsets().equals(sequenceNumbers) && latestSequence.exclusiveStartPartitions.equals( + exclusivePartitions) && !finish) || + (latestSequence.getEndOffsets().equals(sequenceNumbers) && finish)) { + log.warn("Ignoring duplicate request, end sequences already set for sequences [%s]", sequenceNumbers); + return Response.ok(sequenceNumbers).build(); + } else if (latestSequence.isCheckpointed()) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(StringUtils.format( + "WTH?! Sequence [%s] has already endOffsets set, cannot set to [%s]", + latestSequence, + sequenceNumbers + )).build(); + } else if (!isPaused()) { + return Response.status(Response.Status.BAD_REQUEST) + .entity("Task must be paused before changing the end sequences") + .build(); + } + + for (Map.Entry entry : sequenceNumbers.entrySet()) { + if (createSequenceNumber(entry.getValue()).compareTo(createSequenceNumber(currOffsets.get(entry.getKey()))) + < 0) { + return Response.status(Response.Status.BAD_REQUEST) + .entity( + StringUtils.format( + "End sequence must be >= current sequence for partition [%s] (current: %s)", + entry.getKey(), + currOffsets.get(entry.getKey()) + ) + ) + .build(); + } + } + + resetNextCheckpointTime(); + latestSequence.setEndOffsets(sequenceNumbers); + + if (finish) { + log.info("Updating endOffsets from [%s] to [%s]", endOffsets, sequenceNumbers); + endOffsets.putAll(sequenceNumbers); + } else { + exclusiveStartingPartitions.addAll(exclusivePartitions); + + // create new sequence + final SequenceMetadata newSequence = new SequenceMetadata( + latestSequence.getSequenceId() + 1, + StringUtils.format("%s_%d", ioConfig.getBaseSequenceName(), latestSequence.getSequenceId() + 1), + sequenceNumbers, + endOffsets, + false, + exclusivePartitions + ); + sequences.add(newSequence); + initialOffsetsSnapshot.putAll(sequenceNumbers); + } + persistSequences(); + } + catch (Exception e) { + log.error(e, "Unable to set end sequences, dying"); + backgroundThreadException = e; + // should resume to immediately finish kafka index task as failed + resume(); + return Response.status(Response.Status.INTERNAL_SERVER_ERROR) + .entity(Throwables.getStackTraceAsString(e)) + .build(); + } + finally { + pauseLock.unlock(); + } + } + + resume(); + + return Response.ok(sequenceNumbers).build(); + } + + private void resetNextCheckpointTime() + { + nextCheckpointTime = DateTimes.nowUtc().plus(tuningConfig.getIntermediateHandoffPeriod()).getMillis(); + } + + @GET + @Path("/checkpoints") + @Produces(MediaType.APPLICATION_JSON) + public Map> getCheckpointsHTTP( + @Context final HttpServletRequest req + ) + { + authorizationCheck(req, Action.READ); + return getCheckpoints(); + } + + private Map> getCheckpoints() + { + return new TreeMap<>(sequences.stream() + .collect(Collectors.toMap( + SequenceMetadata::getSequenceId, + SequenceMetadata::getStartOffsets + ))); + } + + /** + * Signals the ingestion loop to pause. + * + * @return one of the following Responses: 400 Bad Request if the task has started publishing; 202 Accepted if the + * method has timed out and returned before the task has paused; 200 OK with a map of the current partition sequences + * in the response body if the task successfully paused + */ + @POST + @Path("/pause") + @Produces(MediaType.APPLICATION_JSON) + public Response pauseHTTP( + @Context final HttpServletRequest req + ) throws InterruptedException + { + authorizationCheck(req, Action.WRITE); + return pause(); + } + + @VisibleForTesting + public Response pause() throws InterruptedException + { + if (!(status == Status.PAUSED || status == Status.READING)) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(StringUtils.format("Can't pause, task is not in a pausable state (state: [%s])", status)) + .build(); + } + + pauseLock.lockInterruptibly(); + try { + pauseRequested = true; + + pollRetryLock.lockInterruptibly(); + try { + isAwaitingRetry.signalAll(); + } + finally { + pollRetryLock.unlock(); + } + + if (isPaused()) { + shouldResume.signalAll(); // kick the monitor so it re-awaits with the new pauseMillis + } + + long nanos = TimeUnit.SECONDS.toNanos(2); + while (!isPaused()) { + if (nanos <= 0L) { + return Response.status(Response.Status.ACCEPTED) + .entity("Request accepted but task has not yet paused") + .build(); + } + nanos = hasPaused.awaitNanos(nanos); + } + } + finally { + pauseLock.unlock(); + } + + try { + return Response.ok().entity(toolbox.getObjectMapper().writeValueAsString(getCurrentOffsets())).build(); + } + catch (JsonProcessingException e) { + throw Throwables.propagate(e); + } + } + + @POST + @Path("/resume") + public Response resumeHTTP(@Context final HttpServletRequest req) throws InterruptedException + { + authorizationCheck(req, Action.WRITE); + resume(); + return Response.status(Response.Status.OK).build(); + } + + + @VisibleForTesting + public void resume() throws InterruptedException + { + pauseLock.lockInterruptibly(); + try { + pauseRequested = false; + shouldResume.signalAll(); + + long nanos = TimeUnit.SECONDS.toNanos(5); + while (isPaused()) { + if (nanos <= 0L) { + throw new RuntimeException("Resume command was not accepted within 5 seconds"); + } + nanos = shouldResume.awaitNanos(nanos); + } + } + finally { + pauseLock.unlock(); + } + } + + + @GET + @Path("/time/start") + @Produces(MediaType.APPLICATION_JSON) + public DateTime getStartTime(@Context final HttpServletRequest req) + { + authorizationCheck(req, Action.WRITE); + return startTime; + } + + private class SequenceMetadata + { + private final int sequenceId; + private final String sequenceName; + private final Set exclusiveStartPartitions; + private final Set assignments; + private final boolean sentinel; + private boolean checkpointed; + /** + * Lock for accessing {@link #endOffsets} and {@link #checkpointed}. This lock is required because + * {@link #setEndOffsets)} can be called by both the main thread and the HTTP thread. + */ + private final ReentrantLock lock = new ReentrantLock(); + + final Map startOffsets; + final Map endOffsets; + + @JsonCreator + public SequenceMetadata( + @JsonProperty("sequenceId") int sequenceId, + @JsonProperty("sequenceName") String sequenceName, + @JsonProperty("startOffsets") Map startOffsets, + @JsonProperty("endOffsets") Map endOffsets, + @JsonProperty("checkpointed") boolean checkpointed, + @JsonProperty("exclusiveStartPartitions") Set exclusiveStartPartitions + ) + { + Preconditions.checkNotNull(sequenceName); + Preconditions.checkNotNull(startOffsets); + Preconditions.checkNotNull(endOffsets); + this.sequenceId = sequenceId; + this.sequenceName = sequenceName; + this.startOffsets = ImmutableMap.copyOf(startOffsets); + this.endOffsets = new HashMap<>(endOffsets); + this.assignments = new HashSet<>(startOffsets.keySet()); + this.checkpointed = checkpointed; + this.sentinel = false; + this.exclusiveStartPartitions = exclusiveStartPartitions == null + ? Collections.emptySet() + : exclusiveStartPartitions; + } + + @JsonProperty + public Set getExclusiveStartPartitions() + { + return exclusiveStartPartitions; + } + + @JsonProperty + public int getSequenceId() + { + return sequenceId; + } + + @JsonProperty + public boolean isCheckpointed() + { + lock.lock(); + try { + return checkpointed; + } + finally { + lock.unlock(); + } + } + + @JsonProperty + public String getSequenceName() + { + return sequenceName; + } + + @JsonProperty + public Map getStartOffsets() + { + return startOffsets; + } + + @JsonProperty + public Map getEndOffsets() + { + lock.lock(); + try { + return endOffsets; + } + finally { + lock.unlock(); + } + } + + @JsonProperty + public boolean isSentinel() + { + return sentinel; + } + + void setEndOffsets(Map newEndOffsets) + { + lock.lock(); + try { + endOffsets.putAll(newEndOffsets); + checkpointed = true; + } + finally { + lock.unlock(); + } + } + + void updateAssignments(Map nextPartitionOffset) + { + lock.lock(); + try { + assignments.clear(); + nextPartitionOffset.forEach((key, value) -> { + if (endOffsets.get(key).equals(SeekableStreamPartitions.NO_END_SEQUENCE_NUMBER) + || createSequenceNumber(endOffsets.get(key)).compareTo(createSequenceNumber(nextPartitionOffset.get(key))) + > 0) { + assignments.add(key); + } + }); + } + finally { + lock.unlock(); + } + } + + boolean isOpen() + { + return !assignments.isEmpty(); + } + + boolean canHandle(OrderedPartitionableRecord record) + { + lock.lock(); + try { + final OrderedSequenceNumber partitionEndOffset = createSequenceNumber(endOffsets.get(record.getPartitionId())); + final OrderedSequenceNumber partitionStartOffset = createSequenceNumber(startOffsets.get( + record.getPartitionId())); + final OrderedSequenceNumber recordOffset = createSequenceNumber(record.getSequenceNumber()); + if (!isOpen() || recordOffset == null || partitionEndOffset == null || partitionStartOffset == null) { + return false; + } + boolean ret; + if (isStartingSequenceOffsetsExclusive()) { + ret = recordOffset.compareTo(partitionStartOffset) + >= (getExclusiveStartPartitions().contains(record.getPartitionId()) ? 1 : 0); + } else { + ret = recordOffset.compareTo(partitionStartOffset) >= 0; + } + + if (isEndSequenceOffsetsExclusive()) { + ret &= recordOffset.compareTo(partitionEndOffset) < 0; + } else { + ret &= recordOffset.compareTo(partitionEndOffset) <= 0; + } + + return ret; + } + finally { + lock.unlock(); + } + } + + @Override + public String toString() + { + lock.lock(); + try { + return "SequenceMetadata{" + + "sequenceName='" + sequenceName + '\'' + + ", sequenceId=" + sequenceId + + ", startOffsets=" + startOffsets + + ", endOffsets=" + endOffsets + + ", assignments=" + assignments + + ", sentinel=" + sentinel + + ", checkpointed=" + checkpointed + + '}'; + } + finally { + lock.unlock(); + } + } + + Supplier getCommitterSupplier( + String stream, + Map lastPersistedOffsets + ) + { + // Set up committer. + return () -> + new Committer() + { + @Override + public Object getMetadata() + { + lock.lock(); + + try { + Preconditions.checkState( + assignments.isEmpty(), + "This committer can be used only once all the records till sequences [%s] have been consumed, also make" + + " sure to call updateAssignments before using this committer", + endOffsets + ); + + + // merge endOffsets for this sequence with globally lastPersistedOffsets + // This is done because this committer would be persisting only sub set of segments + // corresponding to the current sequence. Generally, lastPersistedOffsets should already + // cover endOffsets but just to be sure take max of sequences and persist that + for (Map.Entry partitionOffset : endOffsets.entrySet()) { + SequenceOffsetType newOffsets = partitionOffset.getValue(); + if (lastPersistedOffsets.containsKey(partitionOffset.getKey()) && + createSequenceNumber(lastPersistedOffsets.get(partitionOffset.getKey())).compareTo( + createSequenceNumber(newOffsets)) > 0) { + newOffsets = lastPersistedOffsets.get(partitionOffset.getKey()); + } + lastPersistedOffsets.put( + partitionOffset.getKey(), + newOffsets + ); + } + + // Publish metadata can be different from persist metadata as we are going to publish only + // subset of segments + return ImmutableMap.of( + METADATA_NEXT_PARTITIONS, new SeekableStreamPartitions<>(stream, lastPersistedOffsets), + METADATA_PUBLISH_PARTITIONS, new SeekableStreamPartitions<>(stream, endOffsets) + ); + } + finally { + lock.unlock(); + } + } + + @Override + public void run() + { + // Do nothing. + } + }; + + } + + TransactionalSegmentPublisher createPublisher(TaskToolbox toolbox, boolean useTransaction) + { + return (segments, commitMetadata) -> { + final SeekableStreamPartitions finalPartitions = deserializeSeekableStreamPartitionsFromMetadata( + toolbox.getObjectMapper(), + ((Map) Preconditions + .checkNotNull(commitMetadata, "commitMetadata")).get(METADATA_PUBLISH_PARTITIONS) + ); + + // Sanity check, we should only be publishing things that match our desired end state. + if (!getEndOffsets().equals(finalPartitions.getPartitionSequenceNumberMap())) { + throw new ISE( + "WTF?! Driver for sequence [%s], attempted to publish invalid metadata[%s].", + toString(), + commitMetadata + ); + } + + final SegmentTransactionalInsertAction action; + + if (useTransaction) { + action = new SegmentTransactionalInsertAction( + segments, + createDataSourceMetadata(new SeekableStreamPartitions<>( + finalPartitions.getStream(), + getStartOffsets() + )), + createDataSourceMetadata(finalPartitions) + ); + } else { + action = new SegmentTransactionalInsertAction(segments, null, null); + } + + log.info("Publishing with isTransaction[%s].", useTransaction); + + return toolbox.getTaskActionClient().submit(action); + }; + } + + } + + private boolean verifyInitialRecordAndSkipExclusivePartition( + final OrderedPartitionableRecord record, + final Map intialSequenceSnapshot + ) + { + if (intialSequenceSnapshot.containsKey(record.getPartitionId())) { + if (!intialSequenceSnapshot.get(record.getPartitionId()).equals(record.getSequenceNumber())) { + throw new ISE( + "Starting sequenceNumber [%s] does not match expected [%s] for partition [%s]", + record.getSequenceNumber(), + intialSequenceSnapshot.get(record.getPartitionId()), + record.getPartitionId() + ); + } + + log.info( + "Verified starting sequenceNumber [%s] for partition [%s]", + record.getSequenceNumber(), record.getPartitionId() + ); + + intialSequenceSnapshot.remove(record.getPartitionId()); + if (intialSequenceSnapshot.isEmpty()) { + log.info("Verified starting sequences for all partitions"); + } + + // check exclusive starting sequence + if (isStartingSequenceOffsetsExclusive() && exclusiveStartingPartitions.contains(record.getPartitionId())) { + log.info("Skipping starting sequenceNumber for partition [%s] marked exclusive", record.getPartitionId()); + + return false; + } + } + + return true; + } + + /** + * deserailizes the checkpoints into of Map> + * + * @param toolbox task toolbox + * @param checkpointsString the json-serialized checkpoint string + * + * @return checkpoint + * + * @throws IOException jsonProcessingException + */ + @Nullable + protected abstract TreeMap> getCheckPointsFromContext( + TaskToolbox toolbox, + String checkpointsString + ) throws IOException; + + /** + * Calculates the sequence number used to update `currentOffsets` after finishing reading a record. + * In Kafka this returns sequenceNumeber + 1 since that's the next expected offset + * In Kinesis this simply returns sequenceNumber, since the sequence numbers in Kinesis are not + * contiguous and finding the next sequence number requires an expensive API call + * + * @param sequenceNumber the sequence number that has already been processed + * + * @return next sequence number to be stored + */ + protected abstract SequenceOffsetType getSequenceNumberToStoreAfterRead(SequenceOffsetType sequenceNumber); + + /** + * deserialzies stored metadata into SeekableStreamPartitions + * + * @param mapper json objectMapper + * @param object metadata + * + * @return SeekableStreamPartitions + */ + protected abstract SeekableStreamPartitions deserializeSeekableStreamPartitionsFromMetadata( + ObjectMapper mapper, + Object object + ); + + /** + * polls the next set of records from the recordSupplier, the main purpose of having a separate method here + * is to catch and handle exceptions specific to Kafka/Kinesis + * + * @param recordSupplier + * @param toolbox + * + * @return list of records polled, can be empty but cannot be null + * + * @throws Exception + */ + @NotNull + protected abstract List> getRecords( + RecordSupplier recordSupplier, + TaskToolbox toolbox + ) throws Exception; + + /** + * creates specific implementations of kafka/kinesis datasource metadata + * + * @param partitions partitions used to create the datasource metadata + * + * @return datasource metadata + */ + protected abstract SeekableStreamDataSourceMetadata createDataSourceMetadata( + SeekableStreamPartitions partitions + ); + + /** + * create a specific implementation of Kafka/Kinesis sequence number/offset used for comparison mostly + * + * @param sequenceNumber + * + * @return a specific OrderedSequenceNumber instance for Kafka/Kinesis + */ + protected abstract OrderedSequenceNumber createSequenceNumber(SequenceOffsetType sequenceNumber); + + /** + * check if the sequence offsets stored in currOffsets are still valid sequence offsets compared to + * earliest sequence offsets fetched from stream + * + * @param toolbox + * @param recordSupplier + * @param assignment + * @param currOffsets + */ + protected abstract void possiblyResetDataSourceMetadata( + TaskToolbox toolbox, + RecordSupplier recordSupplier, + Set> assignment, + Map currOffsets + ); + + /** + * In Kafka, the endOffsets are exclusive, so skip it. + * In Kinesis the endOffsets are inclusive + */ + protected abstract boolean isEndSequenceOffsetsExclusive(); + + /** + * In Kafka, the startingOffsets are inclusive. + * In Kinesis, the startingOffsets are exclusive, except for the first + * partition we read from stream + */ + protected abstract boolean isStartingSequenceOffsetsExclusive(); +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaTuningConfig.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTuningConfig.java similarity index 77% rename from extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaTuningConfig.java rename to indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTuningConfig.java index fd5e666166b3..816949ce4e16 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaTuningConfig.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTuningConfig.java @@ -17,7 +17,7 @@ * under the License. */ -package org.apache.druid.indexing.kafka; +package org.apache.druid.indexing.seekablestream; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; @@ -32,10 +32,11 @@ import java.io.File; import java.util.Objects; -public class KafkaTuningConfig implements TuningConfig, AppenderatorConfig +public abstract class SeekableStreamIndexTaskTuningConfig implements TuningConfig, AppenderatorConfig { private static final int DEFAULT_MAX_ROWS_PER_SEGMENT = 5_000_000; private static final boolean DEFAULT_RESET_OFFSET_AUTOMATICALLY = false; + private static final boolean DEFAULT_SKIP_SEQUENCE_NUMBER_AVAILABILITY_CHECK = false; private final int maxRowsInMemory; private final long maxBytesInMemory; @@ -53,13 +54,14 @@ public class KafkaTuningConfig implements TuningConfig, AppenderatorConfig @Nullable private final SegmentWriteOutMediumFactory segmentWriteOutMediumFactory; private final Period intermediateHandoffPeriod; + private final boolean skipSequenceNumberAvailabilityCheck; private final boolean logParseExceptions; private final int maxParseExceptions; private final int maxSavedParseExceptions; @JsonCreator - public KafkaTuningConfig( + public SeekableStreamIndexTaskTuningConfig( @JsonProperty("maxRowsInMemory") @Nullable Integer maxRowsInMemory, @JsonProperty("maxBytesInMemory") @Nullable Long maxBytesInMemory, @JsonProperty("maxRowsPerSegment") @Nullable Integer maxRowsPerSegment, @@ -73,6 +75,7 @@ public KafkaTuningConfig( @Deprecated @JsonProperty("reportParseExceptions") @Nullable Boolean reportParseExceptions, @JsonProperty("handoffConditionTimeout") @Nullable Long handoffConditionTimeout, @JsonProperty("resetOffsetAutomatically") @Nullable Boolean resetOffsetAutomatically, + @JsonProperty("skipSequenceNumberAvailabilityCheck") Boolean skipSequenceNumberAvailabilityCheck, @JsonProperty("segmentWriteOutMediumFactory") @Nullable SegmentWriteOutMediumFactory segmentWriteOutMediumFactory, @JsonProperty("intermediateHandoffPeriod") @Nullable Period intermediateHandoffPeriod, @JsonProperty("logParseExceptions") @Nullable Boolean logParseExceptions, @@ -85,15 +88,15 @@ public KafkaTuningConfig( this.maxRowsInMemory = maxRowsInMemory == null ? defaults.getMaxRowsInMemory() : maxRowsInMemory; this.maxRowsPerSegment = maxRowsPerSegment == null ? DEFAULT_MAX_ROWS_PER_SEGMENT : maxRowsPerSegment; + this.maxTotalRows = maxTotalRows; // initializing this to 0, it will be lazily initialized to a value // @see server.src.main.java.org.apache.druid.segment.indexing.TuningConfigs#getMaxBytesInMemoryOrDefault(long) this.maxBytesInMemory = maxBytesInMemory == null ? 0 : maxBytesInMemory; - this.maxTotalRows = maxTotalRows; this.intermediatePersistPeriod = intermediatePersistPeriod == null ? defaults.getIntermediatePersistPeriod() : intermediatePersistPeriod; this.basePersistDirectory = defaults.getBasePersistDirectory(); - this.maxPendingPersists = 0; + this.maxPendingPersists = maxPendingPersists == null ? 0 : maxPendingPersists; this.indexSpec = indexSpec == null ? defaults.getIndexSpec() : indexSpec; this.reportParseExceptions = reportParseExceptions == null ? defaults.isReportParseExceptions() @@ -108,40 +111,24 @@ public KafkaTuningConfig( this.intermediateHandoffPeriod = intermediateHandoffPeriod == null ? new Period().withDays(Integer.MAX_VALUE) : intermediateHandoffPeriod; + this.skipSequenceNumberAvailabilityCheck = skipSequenceNumberAvailabilityCheck == null + ? DEFAULT_SKIP_SEQUENCE_NUMBER_AVAILABILITY_CHECK + : skipSequenceNumberAvailabilityCheck; if (this.reportParseExceptions) { this.maxParseExceptions = 0; this.maxSavedParseExceptions = maxSavedParseExceptions == null ? 0 : Math.min(1, maxSavedParseExceptions); } else { - this.maxParseExceptions = maxParseExceptions == null ? TuningConfig.DEFAULT_MAX_PARSE_EXCEPTIONS : maxParseExceptions; + this.maxParseExceptions = maxParseExceptions == null + ? TuningConfig.DEFAULT_MAX_PARSE_EXCEPTIONS + : maxParseExceptions; this.maxSavedParseExceptions = maxSavedParseExceptions == null ? TuningConfig.DEFAULT_MAX_SAVED_PARSE_EXCEPTIONS : maxSavedParseExceptions; } - this.logParseExceptions = logParseExceptions == null ? TuningConfig.DEFAULT_LOG_PARSE_EXCEPTIONS : logParseExceptions; - } - - public static KafkaTuningConfig copyOf(KafkaTuningConfig config) - { - return new KafkaTuningConfig( - config.maxRowsInMemory, - config.maxBytesInMemory, - config.maxRowsPerSegment, - config.maxTotalRows, - config.intermediatePersistPeriod, - config.basePersistDirectory, - config.maxPendingPersists, - config.indexSpec, - true, - config.reportParseExceptions, - config.handoffConditionTimeout, - config.resetOffsetAutomatically, - config.segmentWriteOutMediumFactory, - config.intermediateHandoffPeriod, - config.logParseExceptions, - config.maxParseExceptions, - config.maxSavedParseExceptions - ); + this.logParseExceptions = logParseExceptions == null + ? TuningConfig.DEFAULT_LOG_PARSE_EXCEPTIONS + : logParseExceptions; } @Override @@ -165,7 +152,6 @@ public int getMaxRowsPerSegment() return maxRowsPerSegment; } - @JsonProperty @Override @Nullable @@ -264,29 +250,14 @@ public int getMaxSavedParseExceptions() return maxSavedParseExceptions; } - public KafkaTuningConfig withBasePersistDirectory(File dir) + @JsonProperty + public boolean isSkipSequenceNumberAvailabilityCheck() { - return new KafkaTuningConfig( - maxRowsInMemory, - maxBytesInMemory, - maxRowsPerSegment, - maxTotalRows, - intermediatePersistPeriod, - dir, - maxPendingPersists, - indexSpec, - true, - reportParseExceptions, - handoffConditionTimeout, - resetOffsetAutomatically, - segmentWriteOutMediumFactory, - intermediateHandoffPeriod, - logParseExceptions, - maxParseExceptions, - maxSavedParseExceptions - ); + return skipSequenceNumberAvailabilityCheck; } + public abstract SeekableStreamIndexTaskTuningConfig withBasePersistDirectory(File dir); + @Override public boolean equals(Object o) { @@ -296,23 +267,24 @@ public boolean equals(Object o) if (o == null || getClass() != o.getClass()) { return false; } - KafkaTuningConfig that = (KafkaTuningConfig) o; + SeekableStreamIndexTaskTuningConfig that = (SeekableStreamIndexTaskTuningConfig) o; return maxRowsInMemory == that.maxRowsInMemory && - maxRowsPerSegment == that.maxRowsPerSegment && maxBytesInMemory == that.maxBytesInMemory && - Objects.equals(maxTotalRows, that.maxTotalRows) && + maxRowsPerSegment == that.maxRowsPerSegment && maxPendingPersists == that.maxPendingPersists && reportParseExceptions == that.reportParseExceptions && handoffConditionTimeout == that.handoffConditionTimeout && resetOffsetAutomatically == that.resetOffsetAutomatically && + skipSequenceNumberAvailabilityCheck == that.skipSequenceNumberAvailabilityCheck && + logParseExceptions == that.logParseExceptions && + maxParseExceptions == that.maxParseExceptions && + maxSavedParseExceptions == that.maxSavedParseExceptions && + Objects.equals(maxTotalRows, that.maxTotalRows) && Objects.equals(intermediatePersistPeriod, that.intermediatePersistPeriod) && Objects.equals(basePersistDirectory, that.basePersistDirectory) && Objects.equals(indexSpec, that.indexSpec) && Objects.equals(segmentWriteOutMediumFactory, that.segmentWriteOutMediumFactory) && - Objects.equals(intermediateHandoffPeriod, that.intermediateHandoffPeriod) && - logParseExceptions == that.logParseExceptions && - maxParseExceptions == that.maxParseExceptions && - maxSavedParseExceptions == that.maxSavedParseExceptions; + Objects.equals(intermediateHandoffPeriod, that.intermediateHandoffPeriod); } @Override @@ -320,8 +292,8 @@ public int hashCode() { return Objects.hash( maxRowsInMemory, - maxRowsPerSegment, maxBytesInMemory, + maxRowsPerSegment, maxTotalRows, intermediatePersistPeriod, basePersistDirectory, @@ -332,6 +304,7 @@ public int hashCode() resetOffsetAutomatically, segmentWriteOutMediumFactory, intermediateHandoffPeriod, + skipSequenceNumberAvailabilityCheck, logParseExceptions, maxParseExceptions, maxSavedParseExceptions @@ -339,25 +312,5 @@ public int hashCode() } @Override - public String toString() - { - return "KafkaTuningConfig{" + - "maxRowsInMemory=" + maxRowsInMemory + - ", maxRowsPerSegment=" + maxRowsPerSegment + - ", maxTotalRows=" + maxTotalRows + - ", maxBytesInMemory=" + maxBytesInMemory + - ", intermediatePersistPeriod=" + intermediatePersistPeriod + - ", basePersistDirectory=" + basePersistDirectory + - ", maxPendingPersists=" + maxPendingPersists + - ", indexSpec=" + indexSpec + - ", reportParseExceptions=" + reportParseExceptions + - ", handoffConditionTimeout=" + handoffConditionTimeout + - ", resetOffsetAutomatically=" + resetOffsetAutomatically + - ", segmentWriteOutMediumFactory=" + segmentWriteOutMediumFactory + - ", intermediateHandoffPeriod=" + intermediateHandoffPeriod + - ", logParseExceptions=" + logParseExceptions + - ", maxParseExceptions=" + maxParseExceptions + - ", maxSavedParseExceptions=" + maxSavedParseExceptions + - '}'; - } + public abstract String toString(); } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamPartitions.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamPartitions.java new file mode 100644 index 000000000000..8c034da41fb4 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamPartitions.java @@ -0,0 +1,127 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; + +import javax.validation.constraints.NotNull; +import java.util.Map; +import java.util.Objects; + +/** + * class that encapsulates a partitionIdToSequenceNumberMap of partitionId -> sequenceNumber. + * To be backward compatible with both Kafka and Kinesis datasource metadata when + * deserializing json. Redundant constructor fields stream, topic and + * partitionSequenceNumberMap and partitionOffsetMap are created. Only one of topic, stream + * should have a non-null value and only one of partitionOffsetMap and partitionSequenceNumberMap + * should have a non-null value. + *

+ * Redundant getters + * are used for proper Jackson serialization/deserialization when processing terminologies + * used by Kafka and kinesis (i.e. topic vs. stream) + * + * @param partition id type + * @param sequence number type + */ +public class SeekableStreamPartitions +{ + // this special marker is used by the KinesisSupervisor to set the endOffsets + // of newly created indexing tasks. This is necessary because streaming tasks do not + // have endPartitionOffsets. This marker signals to the task that it should continue + // to ingest data until taskDuration has elapsed or the task was stopped or paused or killed + public static final String NO_END_SEQUENCE_NUMBER = "NO_END_SEQUENCE_NUMBER"; + + // stream/topic + private final String stream; + // partitionId -> sequence number + private final Map partitionIdToSequenceNumberMap; + + @JsonCreator + public SeekableStreamPartitions( + @JsonProperty("stream") final String stream, + // kept for backward compatibility + @JsonProperty("topic") final String topic, + @JsonProperty("partitionSequenceNumberMap") + final Map partitionSequenceNumberMap, + // kept for backward compatibility + @JsonProperty("partitionOffsetMap") final Map partitionOffsetMap + ) + { + this.stream = stream == null ? topic : stream; + this.partitionIdToSequenceNumberMap = ImmutableMap.copyOf(partitionOffsetMap == null + ? partitionSequenceNumberMap + : partitionOffsetMap); + Preconditions.checkArgument(this.stream != null); + Preconditions.checkArgument(partitionIdToSequenceNumberMap != null); + } + + // constructor for backward compatibility + public SeekableStreamPartitions( + @NotNull final String stream, + final Map partitionOffsetMap + ) + { + this(stream, null, partitionOffsetMap, null); + } + + @JsonProperty + public String getStream() + { + return stream; + } + + @JsonProperty + public Map getPartitionSequenceNumberMap() + { + return partitionIdToSequenceNumberMap; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + SeekableStreamPartitions that = (SeekableStreamPartitions) o; + return Objects.equals(stream, that.stream) && + Objects.equals(partitionIdToSequenceNumberMap, that.partitionIdToSequenceNumberMap); + } + + @Override + public int hashCode() + { + return Objects.hash(stream, partitionIdToSequenceNumberMap); + } + + @Override + public String toString() + { + return "SeekableStreamPartitions{" + + "stream/topic='" + stream + '\'' + + ", partitionSequenceNumberMap/partitionOffsetMap=" + partitionIdToSequenceNumberMap + + '}'; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedPartitionableRecord.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedPartitionableRecord.java new file mode 100644 index 000000000000..4dd653e760ff --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedPartitionableRecord.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.common; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import javax.validation.constraints.NotNull; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * Represents a generic record with a PartitionIdType (partition id) and SequenceOffsetType (sequence number) and data + * from a Kafka/Kinesis stream + * + * @param partition id + * @param sequence number + */ +public class OrderedPartitionableRecord +{ + private final String stream; + private final PartitionIdType partitionId; + private final SequenceOffsetType sequenceNumber; + private final List data; + + public OrderedPartitionableRecord( + String stream, + PartitionIdType partitionId, + SequenceOffsetType sequenceNumber, + List data + ) + { + Preconditions.checkNotNull(stream, "stream"); + Preconditions.checkNotNull(partitionId, "partitionId"); + Preconditions.checkNotNull(sequenceNumber, "sequenceNumber"); + this.stream = stream; + this.partitionId = partitionId; + this.sequenceNumber = sequenceNumber; + this.data = data == null ? ImmutableList.of() : data; + } + + public String getStream() + { + return stream; + } + + public PartitionIdType getPartitionId() + { + return partitionId; + } + + public SequenceOffsetType getSequenceNumber() + { + return sequenceNumber; + } + + @NotNull + public List getData() + { + return data; + } + + public StreamPartition getStreamPartition() + { + return StreamPartition.of(stream, partitionId); + } + + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + OrderedPartitionableRecord that = (OrderedPartitionableRecord) o; + + if (data.size() != that.data.size()) { + return false; + } + + for (int i = 0; i < data.size(); i++) { + if (!Arrays.equals(data.get(i), that.data.get(i))) { + return false; + } + } + + return Objects.equals(stream, that.stream) && + Objects.equals(partitionId, that.partitionId) && + Objects.equals(sequenceNumber, that.sequenceNumber); + } + + @Override + public int hashCode() + { + final int hashOfData = data.stream().map(Arrays::hashCode).collect(Collectors.toList()).hashCode(); + return Objects.hash(stream, partitionId, sequenceNumber, hashOfData); + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedSequenceNumber.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedSequenceNumber.java new file mode 100644 index 000000000000..f193488240d9 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/OrderedSequenceNumber.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.common; + + +/** + * Represents a Kafka/Kinesis stream sequence number. Mainly used to do + * comparison and indicate whether the sequence number is exclusive. + *

+ * isExclusive is used to indicate if this sequence number is the starting + * sequence of some Kinesis partition and should be discarded because some + * previous task has already read this sequence number + * + * @param type of sequence number + */ +public abstract class OrderedSequenceNumber + implements Comparable> +{ + private final SequenceOffsetType sequenceNumber; + private final boolean isExclusive; + + protected OrderedSequenceNumber(SequenceOffsetType sequenceNumber, boolean isExclusive) + { + this.sequenceNumber = sequenceNumber; + this.isExclusive = isExclusive; + } + + public SequenceOffsetType get() + { + return sequenceNumber; + } + + public boolean isExclusive() + { + return isExclusive; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/RecordSupplier.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/RecordSupplier.java new file mode 100644 index 000000000000..d9e599da0c80 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/RecordSupplier.java @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.common; + +import com.google.common.annotations.Beta; + +import javax.annotation.Nullable; +import javax.validation.constraints.NotNull; +import java.io.Closeable; +import java.util.Collection; +import java.util.List; +import java.util.Set; + +/** + * The RecordSupplier interface is a wrapper for the incoming seekable data stream + * (i.e. Kafka consumer, Kinesis streams) + * + * @param Partition Number Type + * @param Sequence Number Type + */ +@Beta +public interface RecordSupplier extends Closeable +{ + /** + * assigns the given partitions to this RecordSupplier + * and seek to the earliest sequence number. Previously + * assigned partitions will be replaced. + * + * @param partitions parititions to assign + */ + void assign(Set> partitions); + + /** + * seek to specified sequence number within a specific partition + * + * @param partition partition to seek + * @param sequenceNumber sequence number to seek to + */ + void seek(StreamPartition partition, SequenceOffsetType sequenceNumber) throws InterruptedException; + + /** + * seek a set of partitions to the earliest record position available in the stream + * + * @param partitions partitions to seek + */ + void seekToEarliest(Set> partitions) throws InterruptedException; + + /** + * seek a set of partitions to the latest/newest record position available in the stream + * + * @param partitions partitions to seek + */ + void seekToLatest(Set> partitions) throws InterruptedException; + + /** + * get the current assignment + * + * @return set of assignments + */ + Collection> getAssignment(); + + /** + * poll the record at the current seeked to sequence in stream + * + * @param timeout timeout in milliseconds + * + * @return record + */ + @NotNull + List> poll(long timeout); + + /** + * get the latest sequence number in stream + * + * @param partition target partition + * + * @return latest sequence number + */ + @Nullable + SequenceOffsetType getLatestSequenceNumber(StreamPartition partition); + + /** + * get the earliest sequence number in stream + * + * @param partition target partition + * + * @return earliest sequence number + */ + @Nullable + SequenceOffsetType getEarliestSequenceNumber(StreamPartition partition); + + + /** + * returns the sequence number of the next record + * + * @param partition target partition + * + * @return sequence number + */ + SequenceOffsetType getPosition(StreamPartition partition); + + /** + * returns the set of partitions under the given stream + * + * @param stream name of stream + * + * @return set of partitions + */ + Set getPartitionIds(String stream); + + /** + * close the RecordSupplier + */ + @Override + void close(); +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/StreamPartition.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/StreamPartition.java new file mode 100644 index 000000000000..5666cdb59fb3 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/common/StreamPartition.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.common; + +/** + * Reprents a Kinesis/Kafka partition with stream name and partitionId, + * mostly used by {@link RecordSupplier} and + * {@link org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor} + * + * @param partition id type + */ +public class StreamPartition +{ + private final String stream; + private final PartitionIdType partitionId; + + public StreamPartition(String stream, PartitionIdType partitionId) + { + this.stream = stream; + this.partitionId = partitionId; + } + + public static StreamPartition of(String stream, PartitionType partitionId) + { + return new StreamPartition<>(stream, partitionId); + } + + public String getStream() + { + return stream; + } + + public PartitionIdType getPartitionId() + { + return partitionId; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + StreamPartition that = (StreamPartition) o; + + if (stream != null ? !stream.equals(that.stream) : that.stream != null) { + return false; + } + return !(partitionId != null ? !partitionId.equals(that.partitionId) : that.partitionId != null); + } + + @Override + public int hashCode() + { + int result = stream != null ? stream.hashCode() : 0; + result = 31 * result + (partitionId != null ? partitionId.hashCode() : 0); + return result; + } + + @Override + public String toString() + { + return "StreamPartition{" + + "stream='" + stream + '\'' + + ", partitionId='" + partitionId + '\'' + + '}'; + } +} + diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java new file mode 100644 index 000000000000..ddd854d3bdad --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -0,0 +1,2816 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.supervisor; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Function; +import com.google.common.base.Joiner; +import com.google.common.base.Optional; +import com.google.common.base.Preconditions; +import com.google.common.base.Predicate; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.MoreExecutors; +import org.apache.commons.codec.digest.DigestUtils; +import org.apache.druid.indexer.TaskLocation; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexing.common.IndexTaskClient; +import org.apache.druid.indexing.common.TaskInfoProvider; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.DataSourceMetadata; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskQueue; +import org.apache.druid.indexing.overlord.TaskRunner; +import org.apache.druid.indexing.overlord.TaskRunnerListener; +import org.apache.druid.indexing.overlord.TaskRunnerWorkItem; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.supervisor.Supervisor; +import org.apache.druid.indexing.overlord.supervisor.SupervisorReport; +import org.apache.druid.indexing.seekablestream.SeekableStreamDataSourceMetadata; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClient; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskIOConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.apache.druid.indexing.seekablestream.SeekableStreamPartitions; +import org.apache.druid.indexing.seekablestream.common.OrderedSequenceNumber; +import org.apache.druid.indexing.seekablestream.common.RecordSupplier; +import org.apache.druid.indexing.seekablestream.common.StreamPartition; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.IAE; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Pair; +import org.apache.druid.java.util.common.RetryUtils; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.metadata.EntryExistsException; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import javax.validation.constraints.NotNull; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; +import java.util.SortedMap; +import java.util.TreeMap; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +/** + * this class is the parent class of both the Kafka and Kinesis supervisor. All the main run loop + * logic are similar enough so they're grouped together into this class. + *

+ * Supervisor responsible for managing the SeekableStreamIndexTasks (Kafka/Kinesis) for a single dataSource. At a high level, the class accepts a + * {@link SeekableStreamSupervisorSpec} which includes the stream name (topic / stream) and configuration as well as an ingestion spec which will + * be used to generate the indexing tasks. The run loop periodically refreshes its view of the stream's partitions + * and the list of running indexing tasks and ensures that all partitions are being read from and that there are enough + * tasks to satisfy the desired number of replicas. As tasks complete, new tasks are queued to process the next range of + * stream sequences. + *

+ * + * @param the type of the partition id, for example, partitions in Kafka are int type while partitions in Kinesis are String type + * @param the type of the sequence number or offsets, for example, Kafka uses long offsets while Kinesis uses String sequence numbers + */ +public abstract class SeekableStreamSupervisor + implements Supervisor +{ + public static final String IS_INCREMENTAL_HANDOFF_SUPPORTED = "IS_INCREMENTAL_HANDOFF_SUPPORTED"; + + private static final long MAX_RUN_FREQUENCY_MILLIS = 1000; + private static final long MINIMUM_FUTURE_TIMEOUT_IN_SECONDS = 120; + private static final int MAX_INITIALIZATION_RETRIES = 20; + + private static final EmittingLogger log = new EmittingLogger(SeekableStreamSupervisor.class); + + // Internal data structures + // -------------------------------------------------------- + + /** + * A TaskGroup is the main data structure used by SeekableStreamSupervisor to organize and monitor stream partitions and + * indexing tasks. All the tasks in a TaskGroup should always be doing the same thing (reading the same partitions and + * starting from the same sequences) and if [replicas] is configured to be 1, a TaskGroup will contain a single task (the + * exception being if the supervisor started up and discovered and adopted some already running tasks). At any given + * time, there should only be up to a maximum of [taskCount] actively-reading task groups (tracked in the [activelyReadingTaskGroups] + * map) + zero or more pending-completion task groups (tracked in [pendingCompletionTaskGroups]). + */ + private class TaskGroup + { + final int groupId; + + // This specifies the partitions and starting sequences for this task group. It is set on group creation from the data + // in [partitionGroups] and never changes during the lifetime of this task group, which will live until a task in + // this task group has completed successfully, at which point this will be destroyed and a new task group will be + // created with new starting sequences. This allows us to create replacement tasks for failed tasks that process the + // same sequences, even if the values in [partitionGroups] has been changed. + final ImmutableMap startingSequences; + + final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + final Optional minimumMessageTime; + final Optional maximumMessageTime; + final Set exclusiveStartSequenceNumberPartitions; + final TreeMap> checkpointSequences = new TreeMap<>(); + final String baseSequenceName; + DateTime completionTimeout; // is set after signalTasksToFinish(); if not done by timeout, take corrective action + + TaskGroup( + int groupId, + ImmutableMap startingSequences, + Optional minimumMessageTime, + Optional maximumMessageTime, + Set exclusiveStartSequenceNumberPartitions + ) + { + this.groupId = groupId; + this.startingSequences = startingSequences; + this.minimumMessageTime = minimumMessageTime; + this.maximumMessageTime = maximumMessageTime; + this.checkpointSequences.put(0, startingSequences); + this.exclusiveStartSequenceNumberPartitions = exclusiveStartSequenceNumberPartitions != null + ? exclusiveStartSequenceNumberPartitions + : new HashSet<>(); + this.baseSequenceName = generateSequenceName(startingSequences, minimumMessageTime, maximumMessageTime); + } + + int addNewCheckpoint(Map checkpoint) + { + checkpointSequences.put(checkpointSequences.lastKey() + 1, checkpoint); + return checkpointSequences.lastKey(); + } + + Set taskIds() + { + return tasks.keySet(); + } + + } + + private class TaskData + { + volatile TaskStatus status; + volatile DateTime startTime; + volatile Map currentSequences = new HashMap<>(); + + @Override + public String toString() + { + return "TaskData{" + + "status=" + status + + ", startTime=" + startTime + + ", checkpointSequences=" + currentSequences + + '}'; + } + } + + /** + * Notice is used to queue tasks that are internal to the supervisor + */ + private interface Notice + { + void handle() throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException; + } + + private static class StatsFromTaskResult + { + private final String groupId; + private final String taskId; + private final Map stats; + + public StatsFromTaskResult( + int groupId, + String taskId, + Map stats + ) + { + this.groupId = String.valueOf(groupId); + this.taskId = taskId; + this.stats = stats; + } + + public String getGroupId() + { + return groupId; + } + + public String getTaskId() + { + return taskId; + } + + public Map getStats() + { + return stats; + } + } + + + private class RunNotice implements Notice + { + @Override + public void handle() throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException + { + long nowTime = System.currentTimeMillis(); + if (nowTime - lastRunTime < MAX_RUN_FREQUENCY_MILLIS) { + return; + } + lastRunTime = nowTime; + + runInternal(); + } + } + + private class GracefulShutdownNotice extends ShutdownNotice + { + @Override + public void handle() throws InterruptedException, ExecutionException, TimeoutException + { + gracefulShutdownInternal(); + super.handle(); + } + } + + private class ShutdownNotice implements Notice + { + @Override + public void handle() throws InterruptedException, ExecutionException, TimeoutException + { + recordSupplier.close(); + + synchronized (stopLock) { + stopped = true; + stopLock.notifyAll(); + } + } + } + + private class ResetNotice implements Notice + { + final DataSourceMetadata dataSourceMetadata; + + ResetNotice(DataSourceMetadata dataSourceMetadata) + { + this.dataSourceMetadata = dataSourceMetadata; + } + + @Override + public void handle() + { + resetInternal(dataSourceMetadata); + } + } + + protected class CheckpointNotice implements Notice + { + @Nullable + private final Integer nullableTaskGroupId; + @Deprecated + private final String baseSequenceName; + private final SeekableStreamDataSourceMetadata previousCheckpoint; + private final SeekableStreamDataSourceMetadata currentCheckpoint; + + public CheckpointNotice( + @Nullable Integer nullableTaskGroupId, + @Deprecated String baseSequenceName, + SeekableStreamDataSourceMetadata previousCheckpoint, + SeekableStreamDataSourceMetadata currentCheckpoint + ) + { + this.baseSequenceName = baseSequenceName; + this.nullableTaskGroupId = nullableTaskGroupId; + this.previousCheckpoint = previousCheckpoint; + this.currentCheckpoint = currentCheckpoint; + } + + @Override + public void handle() throws ExecutionException, InterruptedException + { + // Find taskGroupId using taskId if it's null. It can be null while rolling update. + final int taskGroupId; + if (nullableTaskGroupId == null) { + // We search taskId in activelyReadingTaskGroups and pendingCompletionTaskGroups sequentially. This should be fine because + // 1) a taskGroup can be moved from activelyReadingTaskGroups to pendingCompletionTaskGroups in RunNotice + // (see checkTaskDuration()). + // 2) Notices are proceesed by a single thread. So, CheckpointNotice and RunNotice cannot be processed at the + // same time. + final java.util.Optional maybeGroupId = activelyReadingTaskGroups + .entrySet() + .stream() + .filter(entry -> { + final TaskGroup taskGroup = entry.getValue(); + return taskGroup.baseSequenceName.equals(baseSequenceName); + }) + .findAny() + .map(Entry::getKey); + + taskGroupId = maybeGroupId.orElseGet(() -> pendingCompletionTaskGroups + .entrySet() + .stream() + .filter(entry -> { + final List taskGroups = entry.getValue(); + return taskGroups.stream().anyMatch(group -> group.baseSequenceName.equals(baseSequenceName)); + }) + .findAny() + .orElseThrow(() -> new ISE("Cannot find taskGroup for baseSequenceName[%s]", baseSequenceName)) + .getKey()); + + } else { + taskGroupId = nullableTaskGroupId; + } + + // check for consistency + // if already received request for this sequenceName and dataSourceMetadata combination then return + final TaskGroup taskGroup = activelyReadingTaskGroups.get(taskGroupId); + + if (isValidTaskGroup(taskGroupId, taskGroup)) { + final TreeMap> checkpoints = taskGroup.checkpointSequences; + + // check validity of previousCheckpoint + int index = checkpoints.size(); + for (int sequenceId : checkpoints.descendingKeySet()) { + Map checkpoint = checkpoints.get(sequenceId); + // We have already verified the stream of the current checkpoint is same with that in ioConfig. + // See checkpoint(). + if (checkpoint.equals(previousCheckpoint.getSeekableStreamPartitions() + .getPartitionSequenceNumberMap() + )) { + break; + } + index--; + } + if (index == 0) { + throw new ISE("No such previous checkpoint [%s] found", previousCheckpoint); + } else if (index < checkpoints.size()) { + // if the found checkpoint is not the latest one then already checkpointed by a replica + Preconditions.checkState(index == checkpoints.size() - 1, "checkpoint consistency failure"); + log.info("Already checkpointed with sequences [%s]", checkpoints.lastEntry().getValue()); + return; + } + final Map newCheckpoint = checkpointTaskGroup(taskGroup, false).get(); + taskGroup.addNewCheckpoint(newCheckpoint); + log.info("Handled checkpoint notice, new checkpoint is [%s] for taskGroup [%s]", newCheckpoint, taskGroupId); + } + } + + boolean isValidTaskGroup(int taskGroupId, @Nullable TaskGroup taskGroup) + { + if (taskGroup == null) { + // taskGroup might be in pendingCompletionTaskGroups or partitionGroups + if (pendingCompletionTaskGroups.containsKey(taskGroupId)) { + log.warn( + "Ignoring checkpoint request because taskGroup[%d] has already stopped indexing and is waiting for " + + "publishing segments", + taskGroupId + ); + return false; + } else if (partitionGroups.containsKey(taskGroupId)) { + log.warn("Ignoring checkpoint request because taskGroup[%d] is inactive", taskGroupId); + return false; + } else { + throw new ISE("WTH?! cannot find taskGroup [%s] among all activelyReadingTaskGroups [%s]", taskGroupId, + activelyReadingTaskGroups + ); + } + } + + return true; + } + } + + + // Map<{group RandomIdUtils}, {actively reading task group}>; see documentation for TaskGroup class + private final ConcurrentHashMap activelyReadingTaskGroups = new ConcurrentHashMap<>(); + + // After telling a taskGroup to stop reading and begin publishing a segment, it is moved from [activelyReadingTaskGroups] to here so + // we can monitor its status while we queue new tasks to read the next range of sequences. This is a list since we could + // have multiple sets of tasks publishing at once if time-to-publish > taskDuration. + // Map<{group RandomIdUtils}, List<{pending completion task groups}>> + private final ConcurrentHashMap> pendingCompletionTaskGroups = new ConcurrentHashMap<>(); + + // The starting sequence for a new partition in [partitionGroups] is initially set to getNotSetMarker(). When a new task group + // is created and is assigned partitions, if the sequence in [partitionGroups] is getNotSetMarker() it will take the starting + // sequence value from the metadata store, and if it can't find it there, from stream. Once a task begins + // publishing, the sequence in partitionGroups will be updated to the ending sequence of the publishing-but-not-yet- + // completed task, which will cause the next set of tasks to begin reading from where the previous task left + // off. If that previous task now fails, we will set the sequence in [partitionGroups] back to getNotSetMarker() which will + // cause successive tasks to again grab their starting sequence from metadata store. This mechanism allows us to + // start up successive tasks without waiting for the previous tasks to succeed and still be able to handle task + // failures during publishing. + // Map<{group RandomIdUtils}, Map<{partition RandomIdUtils}, {startingOffset}>> + private final ConcurrentHashMap> partitionGroups = new ConcurrentHashMap<>(); + + protected final ObjectMapper sortingMapper; + protected final List partitionIds = new CopyOnWriteArrayList<>(); + protected volatile DateTime sequenceLastUpdated; + + + private final Set subsequentlyDiscoveredPartitions = new HashSet<>(); + private final TaskStorage taskStorage; + private final TaskMaster taskMaster; + private final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator; + private final SeekableStreamIndexTaskClient taskClient; + private final SeekableStreamSupervisorSpec spec; + private final String dataSource; + private final SeekableStreamSupervisorIOConfig ioConfig; + private final SeekableStreamSupervisorTuningConfig tuningConfig; + private final SeekableStreamIndexTaskTuningConfig taskTuningConfig; + private final String supervisorId; + private final TaskInfoProvider taskInfoProvider; + private final long futureTimeoutInSeconds; // how long to wait for async operations to complete + private final RowIngestionMetersFactory rowIngestionMetersFactory; + private final ExecutorService exec; + private final ScheduledExecutorService scheduledExec; + private final ScheduledExecutorService reportingExec; + private final ListeningExecutorService workerExec; + private final BlockingQueue notices = new LinkedBlockingDeque<>(); + private final Object stopLock = new Object(); + private final Object stateChangeLock = new Object(); + private final Object recordSupplierLock = new Object(); + + private final boolean useExclusiveStartingSequence; + private boolean listenerRegistered = false; + private long lastRunTime; + private int initRetryCounter = 0; + private volatile DateTime firstRunTime; + private volatile RecordSupplier recordSupplier; + private volatile boolean started = false; + private volatile boolean stopped = false; + private volatile boolean lifecycleStarted = false; + + + public SeekableStreamSupervisor( + final String supervisorId, + final TaskStorage taskStorage, + final TaskMaster taskMaster, + final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + final SeekableStreamIndexTaskClientFactory> taskClientFactory, + final ObjectMapper mapper, + final SeekableStreamSupervisorSpec spec, + final RowIngestionMetersFactory rowIngestionMetersFactory, + final boolean useExclusiveStartingSequence + ) + { + this.taskStorage = taskStorage; + this.taskMaster = taskMaster; + this.indexerMetadataStorageCoordinator = indexerMetadataStorageCoordinator; + this.sortingMapper = mapper.copy().configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true); + this.spec = spec; + this.rowIngestionMetersFactory = rowIngestionMetersFactory; + this.useExclusiveStartingSequence = useExclusiveStartingSequence; + + this.dataSource = spec.getDataSchema().getDataSource(); + this.ioConfig = spec.getIoConfig(); + this.tuningConfig = spec.getTuningConfig(); + this.taskTuningConfig = this.tuningConfig.convertToTaskTuningConfig(); + this.supervisorId = supervisorId; + this.exec = Execs.singleThreaded(supervisorId); + this.scheduledExec = Execs.scheduledSingleThreaded(supervisorId + "-Scheduler-%d"); + this.reportingExec = Execs.scheduledSingleThreaded(supervisorId + "-Reporting-%d"); + + int workerThreads = (this.tuningConfig.getWorkerThreads() != null + ? this.tuningConfig.getWorkerThreads() + : Math.min(10, this.ioConfig.getTaskCount())); + + this.workerExec = MoreExecutors.listeningDecorator(Execs.multiThreaded(workerThreads, supervisorId + "-Worker-%d")); + log.info("Created worker pool with [%d] threads for dataSource [%s]", workerThreads, this.dataSource); + + this.taskInfoProvider = new TaskInfoProvider() + { + @Override + public TaskLocation getTaskLocation(final String id) + { + Preconditions.checkNotNull(id, "id"); + Optional taskRunner = taskMaster.getTaskRunner(); + if (taskRunner.isPresent()) { + Optional item = Iterables.tryFind( + taskRunner.get().getRunningTasks(), + (Predicate) taskRunnerWorkItem -> id.equals(taskRunnerWorkItem.getTaskId()) + ); + + if (item.isPresent()) { + return item.get().getLocation(); + } + } else { + log.error("Failed to get task runner because I'm not the leader!"); + } + + return TaskLocation.unknown(); + } + + @Override + public Optional getTaskStatus(String id) + { + return taskStorage.getStatus(id); + } + }; + + this.futureTimeoutInSeconds = Math.max( + MINIMUM_FUTURE_TIMEOUT_IN_SECONDS, + tuningConfig.getChatRetries() * (tuningConfig.getHttpTimeout().getStandardSeconds() + + IndexTaskClient.MAX_RETRY_WAIT_SECONDS) + ); + + int chatThreads = (this.tuningConfig.getChatThreads() != null + ? this.tuningConfig.getChatThreads() + : Math.min(10, this.ioConfig.getTaskCount() * this.ioConfig.getReplicas())); + this.taskClient = taskClientFactory.build( + taskInfoProvider, + dataSource, + chatThreads, + this.tuningConfig.getHttpTimeout(), + this.tuningConfig.getChatRetries() + ); + log.info( + "Created taskClient with dataSource[%s] chatThreads[%d] httpTimeout[%s] chatRetries[%d]", + dataSource, + chatThreads, + this.tuningConfig.getHttpTimeout(), + this.tuningConfig.getChatRetries() + ); + } + + @Override + public void start() + { + synchronized (stateChangeLock) { + Preconditions.checkState(!lifecycleStarted, "already started"); + Preconditions.checkState(!exec.isShutdown(), "already stopped"); + + // Try normal initialization first, if that fails then schedule periodic initialization retries + try { + tryInit(); + } + catch (Exception e) { + if (!started) { + log.warn( + "First initialization attempt failed for SeekableStreamSupervisor[%s], starting retries...", + dataSource + ); + + exec.submit( + () -> { + try { + RetryUtils.retry( + () -> { + tryInit(); + return 0; + }, + (throwable) -> !started, + 0, + MAX_INITIALIZATION_RETRIES, + null, + null + ); + } + catch (Exception e2) { + log.makeAlert( + "Failed to initialize after %s retries, aborting. Please resubmit the supervisor spec to restart this supervisor [%s]", + MAX_INITIALIZATION_RETRIES, + supervisorId + ).emit(); + throw new RuntimeException(e2); + } + } + ); + } + } + lifecycleStarted = true; + } + } + + @Override + public void stop(boolean stopGracefully) + { + synchronized (stateChangeLock) { + Preconditions.checkState(lifecycleStarted, "lifecycle not started"); + + log.info("Beginning shutdown of [%s]", supervisorId); + + try { + scheduledExec.shutdownNow(); // stop recurring executions + reportingExec.shutdownNow(); + recordSupplier.close(); + + if (started) { + Optional taskRunner = taskMaster.getTaskRunner(); + if (taskRunner.isPresent()) { + taskRunner.get().unregisterListener(supervisorId); + } + + // Stopping gracefully will synchronize the end sequences of the tasks and signal them to publish, and will block + // until the tasks have acknowledged or timed out. We want this behavior when we're explicitly shut down through + // the API, but if we shut down for other reasons (e.g. we lose leadership) we want to just stop and leave the + // tasks as they are. + synchronized (stopLock) { + if (stopGracefully) { + log.info("Posting GracefulShutdownNotice, signalling managed tasks to complete and publish"); + notices.add(new GracefulShutdownNotice()); + } else { + log.info("Posting ShutdownNotice"); + notices.add(new ShutdownNotice()); + } + + long shutdownTimeoutMillis = tuningConfig.getShutdownTimeout().getMillis(); + long endTime = System.currentTimeMillis() + shutdownTimeoutMillis; + while (!stopped) { + long sleepTime = endTime - System.currentTimeMillis(); + if (sleepTime <= 0) { + log.info("Timed out while waiting for shutdown (timeout [%,dms])", shutdownTimeoutMillis); + stopped = true; + break; + } + stopLock.wait(sleepTime); + } + } + log.info("Shutdown notice handled"); + } + + taskClient.close(); + workerExec.shutdownNow(); + exec.shutdownNow(); + started = false; + + log.info("[%s] has stopped", supervisorId); + } + catch (Exception e) { + log.makeAlert(e, "Exception stopping [%s]", supervisorId) + .emit(); + } + } + } + + @Override + public void reset(DataSourceMetadata dataSourceMetadata) + { + log.info("Posting ResetNotice"); + notices.add(new ResetNotice(dataSourceMetadata)); + } + + + @VisibleForTesting + public void tryInit() + { + synchronized (stateChangeLock) { + if (started) { + log.warn("Supervisor was already started, skipping init"); + return; + } + + if (stopped) { + log.warn("Supervisor was already stopped, skipping init."); + return; + } + + try { + recordSupplier = setupRecordSupplier(); + + exec.submit( + () -> { + try { + long pollTimeout = Math.max(ioConfig.getPeriod().getMillis(), MAX_RUN_FREQUENCY_MILLIS); + while (!Thread.currentThread().isInterrupted() && !stopped) { + final Notice notice = notices.poll(pollTimeout, TimeUnit.MILLISECONDS); + if (notice == null) { + continue; + } + + try { + notice.handle(); + } + catch (Throwable e) { + log.makeAlert(e, "SeekableStreamSupervisor[%s] failed to handle notice", dataSource) + .addData("noticeClass", notice.getClass().getSimpleName()) + .emit(); + } + } + } + catch (InterruptedException e) { + log.info("SeekableStreamSupervisor[%s] interrupted, exiting", dataSource); + } + } + ); + firstRunTime = DateTimes.nowUtc().plus(ioConfig.getStartDelay()); + scheduledExec.scheduleAtFixedRate( + buildRunTask(), + ioConfig.getStartDelay().getMillis(), + Math.max(ioConfig.getPeriod().getMillis(), MAX_RUN_FREQUENCY_MILLIS), + TimeUnit.MILLISECONDS + ); + + scheduleReporting(reportingExec); + + started = true; + log.info( + "Started SeekableStreamSupervisor[%s], first run in [%s], with spec: [%s]", + dataSource, + ioConfig.getStartDelay(), + spec.toString() + ); + } + catch (Exception e) { + if (recordSupplier != null) { + recordSupplier.close(); + } + initRetryCounter++; + log.makeAlert(e, "Exception starting SeekableStreamSupervisor[%s]", dataSource) + .emit(); + + throw new RuntimeException(e); + } + } + } + + private Runnable buildRunTask() + { + return () -> notices.add(new RunNotice()); + } + + @Override + public SupervisorReport getStatus() + { + return generateReport(true); + } + + private SupervisorReport> generateReport( + boolean includeOffsets + ) + { + int numPartitions = partitionGroups.values().stream().mapToInt(Map::size).sum(); + + final SeekableStreamSupervisorReportPayload payload = createReportPayload( + numPartitions, + includeOffsets + ); + + SupervisorReport> report = new SupervisorReport<>( + dataSource, + DateTimes.nowUtc(), + payload + ); + + List> taskReports = new ArrayList<>(); + + try { + for (TaskGroup taskGroup : activelyReadingTaskGroups.values()) { + for (Entry entry : taskGroup.tasks.entrySet()) { + String taskId = entry.getKey(); + @Nullable + DateTime startTime = entry.getValue().startTime; + Map currentOffsets = entry.getValue().currentSequences; + Long remainingSeconds = null; + if (startTime != null) { + long elapsedMillis = System.currentTimeMillis() - startTime.getMillis(); + long remainingMillis = Math.max(0, ioConfig.getTaskDuration().getMillis() - elapsedMillis); + remainingSeconds = TimeUnit.MILLISECONDS.toSeconds(remainingMillis); + } + + taskReports.add( + new TaskReportData<>( + taskId, + includeOffsets ? taskGroup.startingSequences : null, + includeOffsets ? currentOffsets : null, + startTime, + remainingSeconds, + TaskReportData.TaskType.ACTIVE, + includeOffsets ? getLagPerPartition(currentOffsets) : null + ) + ); + } + } + + for (List taskGroups : pendingCompletionTaskGroups.values()) { + for (TaskGroup taskGroup : taskGroups) { + for (Entry entry : taskGroup.tasks.entrySet()) { + String taskId = entry.getKey(); + @Nullable + DateTime startTime = entry.getValue().startTime; + Map currentOffsets = entry.getValue().currentSequences; + Long remainingSeconds = null; + if (taskGroup.completionTimeout != null) { + remainingSeconds = Math.max(0, taskGroup.completionTimeout.getMillis() - System.currentTimeMillis()) + / 1000; + } + + taskReports.add( + new TaskReportData<>( + taskId, + includeOffsets ? taskGroup.startingSequences : null, + includeOffsets ? currentOffsets : null, + startTime, + remainingSeconds, + TaskReportData.TaskType.PUBLISHING, + null + ) + ); + } + } + } + + taskReports.forEach(payload::addTask); + } + catch (Exception e) { + log.warn(e, "Failed to generate status report"); + } + + return report; + } + + + @Override + public Map> getStats() + { + try { + return getCurrentTotalStats(); + } + catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + log.error(ie, "getStats() interrupted."); + throw new RuntimeException(ie); + } + catch (ExecutionException | TimeoutException eete) { + throw new RuntimeException(eete); + } + } + + /** + * Collect row ingestion stats from all tasks managed by this supervisor. + * + * @return A map of groupId->taskId->task row stats + * + * @throws InterruptedException + * @throws ExecutionException + * @throws TimeoutException + */ + private Map> getCurrentTotalStats() + throws InterruptedException, ExecutionException, TimeoutException + { + Map> allStats = new HashMap<>(); + final List> futures = new ArrayList<>(); + final List> groupAndTaskIds = new ArrayList<>(); + + for (int groupId : activelyReadingTaskGroups.keySet()) { + TaskGroup group = activelyReadingTaskGroups.get(groupId); + for (String taskId : group.taskIds()) { + futures.add( + Futures.transform( + taskClient.getMovingAveragesAsync(taskId), + (Function, StatsFromTaskResult>) (currentStats) -> new StatsFromTaskResult( + groupId, + taskId, + currentStats + ) + ) + ); + groupAndTaskIds.add(new Pair<>(groupId, taskId)); + } + } + + for (int groupId : pendingCompletionTaskGroups.keySet()) { + List pendingGroups = pendingCompletionTaskGroups.get(groupId); + for (TaskGroup pendingGroup : pendingGroups) { + for (String taskId : pendingGroup.taskIds()) { + futures.add( + Futures.transform( + taskClient.getMovingAveragesAsync(taskId), + (Function, StatsFromTaskResult>) (currentStats) -> new StatsFromTaskResult( + groupId, + taskId, + currentStats + ) + ) + ); + groupAndTaskIds.add(new Pair<>(groupId, taskId)); + } + } + } + + List results = Futures.successfulAsList(futures) + .get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int i = 0; i < results.size(); i++) { + StatsFromTaskResult result = results.get(i); + if (result != null) { + Map groupMap = allStats.computeIfAbsent(result.getGroupId(), k -> new HashMap<>()); + groupMap.put(result.getTaskId(), result.getStats()); + } else { + Pair groupAndTaskId = groupAndTaskIds.get(i); + log.error("Failed to get stats for group[%d]-task[%s]", groupAndTaskId.lhs, groupAndTaskId.rhs); + } + } + + return allStats; + } + + + @VisibleForTesting + public void addTaskGroupToActivelyReadingTaskGroup( + int taskGroupId, + ImmutableMap partitionOffsets, + Optional minMsgTime, + Optional maxMsgTime, + Set tasks, + Set exclusiveStartingSequencePartitions + ) + { + TaskGroup group = new TaskGroup( + taskGroupId, + partitionOffsets, + minMsgTime, + maxMsgTime, + exclusiveStartingSequencePartitions + ); + group.tasks.putAll(tasks.stream().collect(Collectors.toMap(x -> x, x -> new TaskData()))); + if (activelyReadingTaskGroups.putIfAbsent(taskGroupId, group) != null) { + throw new ISE( + "trying to add taskGroup with RandomIdUtils [%s] to actively reading task groups, but group already exists.", + taskGroupId + ); + } + } + + @VisibleForTesting + public void addTaskGroupToPendingCompletionTaskGroup( + int taskGroupId, + ImmutableMap partitionOffsets, + Optional minMsgTime, + Optional maxMsgTime, + Set tasks, + Set exclusiveStartingSequencePartitions + ) + { + TaskGroup group = new TaskGroup( + taskGroupId, + partitionOffsets, + minMsgTime, + maxMsgTime, + exclusiveStartingSequencePartitions + ); + group.tasks.putAll(tasks.stream().collect(Collectors.toMap(x -> x, x -> new TaskData()))); + pendingCompletionTaskGroups.computeIfAbsent(taskGroupId, x -> new CopyOnWriteArrayList<>()) + .add(group); + } + + @VisibleForTesting + public void runInternal() + throws ExecutionException, InterruptedException, TimeoutException, JsonProcessingException + { + possiblyRegisterListener(); + updatePartitionDataFromStream(); + discoverTasks(); + updateTaskStatus(); + checkTaskDuration(); + checkPendingCompletionTasks(); + checkCurrentTaskState(); + // if supervisor is not suspended, ensure required tasks are running + // if suspended, ensure tasks have been requested to gracefully stop + if (!spec.isSuspended()) { + log.info("[%s] supervisor is running.", dataSource); + createNewTasks(); + } else { + log.info("[%s] supervisor is suspended.", dataSource); + gracefulShutdownInternal(); + } + + if (log.isDebugEnabled()) { + log.debug(generateReport(true).toString()); + } else { + log.info(generateReport(false).toString()); + } + } + + private void possiblyRegisterListener() + { + if (listenerRegistered) { + return; + } + + Optional taskRunner = taskMaster.getTaskRunner(); + if (taskRunner.isPresent()) { + taskRunner.get().registerListener( + new TaskRunnerListener() + { + @Override + public String getListenerId() + { + return supervisorId; + } + + @Override + public void locationChanged(final String taskId, final TaskLocation newLocation) + { + // do nothing + } + + @Override + public void statusChanged(String taskId, TaskStatus status) + { + notices.add(new RunNotice()); + } + }, MoreExecutors.sameThreadExecutor() + ); + listenerRegistered = true; + } + } + + @VisibleForTesting + public void gracefulShutdownInternal() throws ExecutionException, InterruptedException, TimeoutException + { + for (TaskGroup taskGroup : activelyReadingTaskGroups.values()) { + for (Entry entry : taskGroup.tasks.entrySet()) { + if (taskInfoProvider.getTaskLocation(entry.getKey()).equals(TaskLocation.unknown())) { + killTask(entry.getKey(), "Killing task for graceful shutdown"); + } else { + entry.getValue().startTime = DateTimes.EPOCH; + } + } + } + + checkTaskDuration(); + } + + @VisibleForTesting + public void resetInternal(DataSourceMetadata dataSourceMetadata) + { + if (dataSourceMetadata == null) { + // Reset everything + boolean result = indexerMetadataStorageCoordinator.deleteDataSourceMetadata(dataSource); + log.info("Reset dataSource[%s] - dataSource metadata entry deleted? [%s]", dataSource, result); + activelyReadingTaskGroups.values() + .forEach(group -> killTasksInGroup( + group, + "DataSourceMetadata is not found while reset" + )); + activelyReadingTaskGroups.clear(); + partitionGroups.clear(); + } else { + + if (!checkSourceMetadataMatch(dataSourceMetadata)) { + throw new IAE( + "Datasource metadata instance does not match required, found instance of [%s]", + dataSourceMetadata.getClass() + ); + } + // Reset only the partitions in dataSourceMetadata if it has not been reset yet + @SuppressWarnings("unchecked") + final SeekableStreamDataSourceMetadata resetMetadata = (SeekableStreamDataSourceMetadata) dataSourceMetadata; + + if (resetMetadata.getSeekableStreamPartitions().getStream().equals(ioConfig.getStream())) { + // metadata can be null + final DataSourceMetadata metadata = indexerMetadataStorageCoordinator.getDataSourceMetadata(dataSource); + if (metadata != null && !checkSourceMetadataMatch(metadata)) { + throw new IAE( + "Datasource metadata instance does not match required, found instance of [%s]", + metadata.getClass() + ); + } + + @SuppressWarnings("unchecked") + final SeekableStreamDataSourceMetadata currentMetadata = (SeekableStreamDataSourceMetadata) metadata; + + // defend against consecutive reset requests from replicas + // as well as the case where the metadata store do not have an entry for the reset partitions + boolean doReset = false; + for (Entry resetPartitionOffset : resetMetadata.getSeekableStreamPartitions() + .getPartitionSequenceNumberMap() + .entrySet()) { + final SequenceOffsetType partitionOffsetInMetadataStore = currentMetadata == null + ? null + : currentMetadata.getSeekableStreamPartitions() + .getPartitionSequenceNumberMap() + .get(resetPartitionOffset.getKey()); + final TaskGroup partitionTaskGroup = activelyReadingTaskGroups.get( + getTaskGroupIdForPartition(resetPartitionOffset.getKey()) + ); + final boolean isSameOffset = partitionTaskGroup != null + && partitionTaskGroup.startingSequences.get(resetPartitionOffset.getKey()) + .equals(resetPartitionOffset.getValue()); + if (partitionOffsetInMetadataStore != null || isSameOffset) { + doReset = true; + break; + } + } + + if (!doReset) { + log.info("Ignoring duplicate reset request [%s]", dataSourceMetadata); + return; + } + + boolean metadataUpdateSuccess = false; + if (currentMetadata == null) { + metadataUpdateSuccess = true; + } else { + final DataSourceMetadata newMetadata = currentMetadata.minus(resetMetadata); + try { + metadataUpdateSuccess = indexerMetadataStorageCoordinator.resetDataSourceMetadata(dataSource, newMetadata); + } + catch (IOException e) { + log.error("Resetting DataSourceMetadata failed [%s]", e.getMessage()); + Throwables.propagate(e); + } + } + if (metadataUpdateSuccess) { + resetMetadata.getSeekableStreamPartitions().getPartitionSequenceNumberMap().keySet().forEach(partition -> { + final int groupId = getTaskGroupIdForPartition(partition); + killTaskGroupForPartitions(ImmutableSet.of(partition), "DataSourceMetadata is updated while reset"); + activelyReadingTaskGroups.remove(groupId); + partitionGroups.get(groupId).replaceAll((partitionId, sequence) -> getNotSetMarker()); + }); + } else { + throw new ISE("Unable to reset metadata"); + } + } else { + log.warn( + "Reset metadata stream [%s] and supervisor's stream name [%s] do not match", + resetMetadata.getSeekableStreamPartitions().getStream(), + ioConfig.getStream() + ); + } + } + + + } + + private void killTask(final String id, String reasonFormat, Object... args) + { + Optional taskQueue = taskMaster.getTaskQueue(); + if (taskQueue.isPresent()) { + taskQueue.get().shutdown(id, reasonFormat, args); + } else { + log.error("Failed to get task queue because I'm not the leader!"); + } + } + + private void killTasksInGroup(TaskGroup taskGroup, String reasonFormat, Object... args) + { + if (taskGroup != null) { + for (String taskId : taskGroup.tasks.keySet()) { + killTask(taskId, reasonFormat, args); + } + } + } + + private void killTaskGroupForPartitions(Set partitions, String reasonFormat, Object... args) + { + for (PartitionIdType partition : partitions) { + int taskGroupId = getTaskGroupIdForPartition(partition); + killTasksInGroup(activelyReadingTaskGroups.get(taskGroupId), reasonFormat, args); + partitionGroups.remove(taskGroupId); + activelyReadingTaskGroups.remove(taskGroupId); + } + } + + private boolean isTaskInPendingCompletionGroups(String taskId) + { + for (List taskGroups : pendingCompletionTaskGroups.values()) { + for (TaskGroup taskGroup : taskGroups) { + if (taskGroup.tasks.containsKey(taskId)) { + return true; + } + } + } + return false; + } + + private void discoverTasks() throws ExecutionException, InterruptedException, TimeoutException + { + int taskCount = 0; + List futureTaskIds = new ArrayList<>(); + List> futures = new ArrayList<>(); + List tasks = taskStorage.getActiveTasks(); + final Map taskGroupsToVerify = new HashMap<>(); + + for (Task task : tasks) { + if (!doesTaskTypeMatchSupervisor(task) || !dataSource.equals(task.getDataSource())) { + continue; + } + + taskCount++; + @SuppressWarnings("unchecked") + final SeekableStreamIndexTask seekableStreamIndexTask = (SeekableStreamIndexTask) task; + final String taskId = task.getId(); + + // Determine which task group this task belongs to based on one of the partitions handled by this task. If we + // later determine that this task is actively reading, we will make sure that it matches our current partition + // allocation (getTaskGroupIdForPartition(partition) should return the same value for every partition being read + // by this task) and kill it if it is not compatible. If the task is instead found to be in the publishing + // state, we will permit it to complete even if it doesn't match our current partition allocation to support + // seamless schema migration. + + Iterator it = seekableStreamIndexTask.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap() + .keySet() + .iterator(); + final Integer taskGroupId = (it.hasNext() ? getTaskGroupIdForPartition(it.next()) : null); + + if (taskGroupId != null) { + // check to see if we already know about this task, either in [activelyReadingTaskGroups] or in [pendingCompletionTaskGroups] + // and if not add it to activelyReadingTaskGroups or pendingCompletionTaskGroups (if status = PUBLISHING) + TaskGroup taskGroup = activelyReadingTaskGroups.get(taskGroupId); + + if (!isTaskInPendingCompletionGroups(taskId) && (taskGroup == null || !taskGroup.tasks.containsKey(taskId))) { + futureTaskIds.add(taskId); + futures.add( + Futures.transform( + taskClient.getStatusAsync(taskId), new Function() + { + @Override + public Boolean apply(SeekableStreamIndexTaskRunner.Status status) + { + try { + log.debug("Task [%s], status [%s]", taskId, status); + if (status == SeekableStreamIndexTaskRunner.Status.PUBLISHING) { + seekableStreamIndexTask.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap() + .keySet() + .forEach( + partition -> addDiscoveredTaskToPendingCompletionTaskGroups( + getTaskGroupIdForPartition(partition), + taskId, + seekableStreamIndexTask.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap() + )); + + // update partitionGroups with the publishing task's sequences (if they are greater than what is + // existing) so that the next tasks will start reading from where this task left off + Map publishingTaskEndOffsets = taskClient.getEndOffsets( + taskId); + + for (Entry entry : publishingTaskEndOffsets.entrySet()) { + PartitionIdType partition = entry.getKey(); + SequenceOffsetType sequence = entry.getValue(); + ConcurrentHashMap partitionOffsets = partitionGroups.get( + getTaskGroupIdForPartition(partition) + ); + + boolean succeeded; + do { + succeeded = true; + SequenceOffsetType previousOffset = partitionOffsets.putIfAbsent(partition, sequence); + if (previousOffset != null + && (makeSequenceNumber(previousOffset).compareTo(makeSequenceNumber(sequence))) < 0) { + succeeded = partitionOffsets.replace(partition, previousOffset, sequence); + } + } while (!succeeded); + } + } else { + for (PartitionIdType partition : seekableStreamIndexTask.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap() + .keySet()) { + if (!taskGroupId.equals(getTaskGroupIdForPartition(partition))) { + log.warn( + "Stopping task [%s] which does not match the expected partition allocation", + taskId + ); + try { + stopTask(taskId, false).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + catch (InterruptedException | ExecutionException | TimeoutException e) { + log.warn(e, "Exception while stopping task"); + } + return false; + } + } + // make sure the task's io and tuning configs match with the supervisor config + // if it is current then only create corresponding taskGroup if it does not exist + if (!isTaskCurrent(taskGroupId, taskId)) { + log.info( + "Stopping task [%s] which does not match the expected parameters and ingestion spec", + taskId + ); + try { + stopTask(taskId, false).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + catch (InterruptedException | ExecutionException | TimeoutException e) { + log.warn(e, "Exception while stopping task"); + } + return false; + } else { + final TaskGroup taskGroup = activelyReadingTaskGroups.computeIfAbsent( + taskGroupId, + k -> { + log.info("Creating a new task group for taskGroupId[%d]", taskGroupId); + return new TaskGroup( + taskGroupId, + ImmutableMap.copyOf( + seekableStreamIndexTask.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap() + ), + seekableStreamIndexTask.getIOConfig().getMinimumMessageTime(), + seekableStreamIndexTask.getIOConfig().getMaximumMessageTime(), + seekableStreamIndexTask.getIOConfig().getExclusiveStartSequenceNumberPartitions() + ); + } + ); + taskGroupsToVerify.put(taskGroupId, taskGroup); + final TaskData prevTaskData = taskGroup.tasks.putIfAbsent(taskId, new TaskData()); + if (prevTaskData != null) { + throw new ISE( + "WTH? a taskGroup[%s] already exists for new task[%s]", + prevTaskData, + taskId + ); + } + } + } + return true; + } + catch (Throwable t) { + log.error(t, "Something bad while discovering task [%s]", taskId); + return null; + } + } + + }, workerExec + ) + ); + } + } + } + + + List results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int i = 0; i < results.size(); i++) { + if (results.get(i) == null) { + String taskId = futureTaskIds.get(i); + killTask(taskId, "Task [%s] failed to return status, killing task", taskId); + } + } + log.debug("Found [%d] seekablestream indexing tasks for dataSource [%s]", taskCount, dataSource); + + // make sure the checkpoints are consistent with each other and with the metadata store + + verifyAndMergeCheckpoints(taskGroupsToVerify.values()); + + } + + private void verifyAndMergeCheckpoints(final Collection taskGroupsToVerify) + { + final List> futures = new ArrayList<>(); + for (TaskGroup taskGroup : taskGroupsToVerify) { + futures.add(workerExec.submit(() -> verifyAndMergeCheckpoints(taskGroup))); + } + try { + Futures.allAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + catch (InterruptedException | ExecutionException | TimeoutException e) { + throw new RuntimeException(e); + } + } + + /** + * This method does two things - + * 1. Makes sure the checkpoints information in the taskGroup is consistent with that of the tasks, if not kill + * inconsistent tasks. + * 2. truncates the checkpoints in the taskGroup corresponding to which segments have been published, so that any newly + * created tasks for the taskGroup start indexing from after the latest published sequences. + */ + private void verifyAndMergeCheckpoints(final TaskGroup taskGroup) + { + final int groupId = taskGroup.groupId; + final List>>> taskSequences = new ArrayList<>(); + final List>>> futures = new ArrayList<>(); + final List taskIds = new ArrayList<>(); + + for (String taskId : taskGroup.taskIds()) { + final ListenableFuture>> checkpointsFuture = taskClient.getCheckpointsAsync( + taskId, + true + ); + futures.add(checkpointsFuture); + taskIds.add(taskId); + } + + try { + List>> futuresResult = + Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int i = 0; i < futuresResult.size(); i++) { + final TreeMap> checkpoints = futuresResult.get(i); + final String taskId = taskIds.get(i); + if (checkpoints == null) { + try { + // catch the exception in failed futures + futures.get(i).get(); + } + catch (Exception e) { + log.error(e, "Problem while getting checkpoints for task [%s], killing the task", taskId); + killTask(taskId, "Exception[%s] while getting checkpoints", e.getClass()); + taskGroup.tasks.remove(taskId); + } + } else if (checkpoints.isEmpty()) { + log.warn("Ignoring task [%s], as probably it is not started running yet", taskId); + } else { + taskSequences.add(new Pair<>(taskId, checkpoints)); + } + } + } + catch (Exception e) { + throw new RuntimeException(e); + } + + final DataSourceMetadata rawDataSourceMetadata = indexerMetadataStorageCoordinator.getDataSourceMetadata(dataSource); + + if (rawDataSourceMetadata != null && !checkSourceMetadataMatch(rawDataSourceMetadata)) { + throw new IAE( + "Datasource metadata instance does not match required, found instance of [%s]", + rawDataSourceMetadata.getClass() + ); + } + + @SuppressWarnings("unchecked") + final SeekableStreamDataSourceMetadata latestDataSourceMetadata = (SeekableStreamDataSourceMetadata) rawDataSourceMetadata; + + final boolean hasValidOffsetsFromDb = latestDataSourceMetadata != null && + latestDataSourceMetadata.getSeekableStreamPartitions() != null && + ioConfig.getStream().equals( + latestDataSourceMetadata.getSeekableStreamPartitions().getStream() + ); + final Map latestOffsetsFromDb; + if (hasValidOffsetsFromDb) { + latestOffsetsFromDb = latestDataSourceMetadata.getSeekableStreamPartitions().getPartitionSequenceNumberMap(); + } else { + latestOffsetsFromDb = null; + } + + // order tasks of this taskGroup by the latest sequenceId + taskSequences.sort((o1, o2) -> o2.rhs.firstKey().compareTo(o1.rhs.firstKey())); + + final Set tasksToKill = new HashSet<>(); + final AtomicInteger earliestConsistentSequenceId = new AtomicInteger(-1); + int taskIndex = 0; + + while (taskIndex < taskSequences.size()) { + TreeMap> taskCheckpoints = taskSequences.get(taskIndex).rhs; + String taskId = taskSequences.get(taskIndex).lhs; + if (earliestConsistentSequenceId.get() == -1) { + // find the first replica task with earliest sequenceId consistent with datasource metadata in the metadata + // store + if (taskCheckpoints.entrySet().stream().anyMatch( + sequenceCheckpoint -> sequenceCheckpoint.getValue().entrySet().stream().allMatch( + partitionOffset -> { + OrderedSequenceNumber sequence = makeSequenceNumber(partitionOffset.getValue()); + OrderedSequenceNumber latestOffset = makeSequenceNumber( + latestOffsetsFromDb == null ? partitionOffset.getValue() : + latestOffsetsFromDb.getOrDefault( + partitionOffset + .getKey(), + partitionOffset + .getValue() + ) + ); + + return sequence.compareTo(latestOffset) == 0; + } + ) && earliestConsistentSequenceId.compareAndSet(-1, sequenceCheckpoint.getKey())) || ( + pendingCompletionTaskGroups.getOrDefault(groupId, new CopyOnWriteArrayList<>()).size() > 0 + && earliestConsistentSequenceId.compareAndSet(-1, taskCheckpoints.firstKey()))) { + final SortedMap> latestCheckpoints = new TreeMap<>( + taskCheckpoints.tailMap(earliestConsistentSequenceId.get()) + ); + log.info("Setting taskGroup sequences to [%s] for group [%d]", latestCheckpoints, groupId); + taskGroup.checkpointSequences.clear(); + taskGroup.checkpointSequences.putAll(latestCheckpoints); + } else { + log.debug( + "Adding task [%s] to kill list, checkpoints[%s], latestoffsets from DB [%s]", + taskId, + taskCheckpoints, + latestOffsetsFromDb + ); + tasksToKill.add(taskId); + } + } else { + // check consistency with taskGroup sequences + if (taskCheckpoints.get(taskGroup.checkpointSequences.firstKey()) == null + || !(taskCheckpoints.get(taskGroup.checkpointSequences.firstKey()) + .equals(taskGroup.checkpointSequences.firstEntry().getValue())) + || taskCheckpoints.tailMap(taskGroup.checkpointSequences.firstKey()).size() + != taskGroup.checkpointSequences.size()) { + log.debug( + "Adding task [%s] to kill list, checkpoints[%s], taskgroup checkpoints [%s]", + taskId, + taskCheckpoints, + taskGroup.checkpointSequences + ); + tasksToKill.add(taskId); + } + } + taskIndex++; + } + + if ((tasksToKill.size() > 0 && tasksToKill.size() == taskGroup.tasks.size()) || + (taskGroup.tasks.size() == 0 + && pendingCompletionTaskGroups.getOrDefault(groupId, new CopyOnWriteArrayList<>()).size() == 0)) { + // killing all tasks or no task left in the group ? + // clear state about the taskgroup so that get latest sequence information is fetched from metadata store + log.warn("Clearing task group [%d] information as no valid tasks left the group", groupId); + activelyReadingTaskGroups.remove(groupId); + partitionGroups.get(groupId).replaceAll((partition, sequence) -> getNotSetMarker()); + } + + taskSequences.stream().filter(taskIdSequences -> tasksToKill.contains(taskIdSequences.lhs)).forEach( + sequenceCheckpoint -> { + killTask( + sequenceCheckpoint.lhs, + "Killing task [%s], as its checkpoints [%s] are not consistent with group checkpoints[%s] or latest " + + "persisted sequences in metadata store [%s]", + sequenceCheckpoint.lhs, + sequenceCheckpoint.rhs, + taskGroup.checkpointSequences, + latestOffsetsFromDb + ); + taskGroup.tasks.remove(sequenceCheckpoint.lhs); + } + ); + } + + private void addDiscoveredTaskToPendingCompletionTaskGroups( + int groupId, + String taskId, + Map startingPartitions + ) + { + final CopyOnWriteArrayList taskGroupList = pendingCompletionTaskGroups.computeIfAbsent( + groupId, + k -> new CopyOnWriteArrayList<>() + ); + for (TaskGroup taskGroup : taskGroupList) { + if (taskGroup.startingSequences.equals(startingPartitions)) { + if (taskGroup.tasks.putIfAbsent(taskId, new TaskData()) == null) { + log.info("Added discovered task [%s] to existing pending task group [%s]", taskId, groupId); + } + return; + } + } + + log.info("Creating new pending completion task group [%s] for discovered task [%s]", groupId, taskId); + + // reading the minimumMessageTime & maximumMessageTime from the publishing task and setting it here is not necessary as this task cannot + // change to a state where it will read any more events + TaskGroup newTaskGroup = new TaskGroup( + groupId, + ImmutableMap.copyOf(startingPartitions), + Optional.absent(), + Optional.absent(), + null + ); + + newTaskGroup.tasks.put(taskId, new TaskData()); + newTaskGroup.completionTimeout = DateTimes.nowUtc().plus(ioConfig.getCompletionTimeout()); + + taskGroupList.add(newTaskGroup); + } + + private ListenableFuture stopTask(final String id, final boolean publish) + { + return Futures.transform( + taskClient.stopAsync(id, publish), new Function() + { + @Nullable + @Override + public Void apply(@Nullable Boolean result) + { + if (result == null || !result) { + log.info("Task [%s] failed to stop in a timely manner, killing task", id); + killTask(id, "Task [%s] failed to stop in a timely manner, killing task", id); + } + return null; + } + } + ); + } + + private boolean isTaskCurrent(int taskGroupId, String taskId) + { + Optional taskOptional = taskStorage.getTask(taskId); + if (!taskOptional.isPresent() || !doesTaskTypeMatchSupervisor(taskOptional.get())) { + return false; + } + + @SuppressWarnings("unchecked") + SeekableStreamIndexTask task = (SeekableStreamIndexTask) taskOptional + .get(); + + String taskSequenceName = task.getIOConfig().getBaseSequenceName(); + if (activelyReadingTaskGroups.get(taskGroupId) != null) { + return Preconditions + .checkNotNull(activelyReadingTaskGroups.get(taskGroupId), "null taskGroup for taskId[%s]", taskGroupId) + .baseSequenceName + .equals(taskSequenceName); + } else { + return generateSequenceName( + task.getIOConfig() + .getStartPartitions() + .getPartitionSequenceNumberMap(), + task.getIOConfig().getMinimumMessageTime(), + task.getIOConfig().getMaximumMessageTime() + ).equals(taskSequenceName); + } + } + + @VisibleForTesting + protected String generateSequenceName( + Map startPartitions, + Optional minimumMessageTime, + Optional maximumMessageTime + ) + { + StringBuilder sb = new StringBuilder(); + + for (Entry entry : startPartitions.entrySet()) { + sb.append(StringUtils.format("+%s(%s)", entry.getKey().toString(), entry.getValue().toString())); + } + String partitionOffsetStr = sb.toString().substring(1); + + String minMsgTimeStr = (minimumMessageTime.isPresent() ? String.valueOf(minimumMessageTime.get().getMillis()) : ""); + String maxMsgTimeStr = (maximumMessageTime.isPresent() ? String.valueOf(maximumMessageTime.get().getMillis()) : ""); + + String dataSchema, tuningConfig; + try { + dataSchema = sortingMapper.writeValueAsString(spec.getDataSchema()); + tuningConfig = sortingMapper.writeValueAsString(taskTuningConfig); + } + catch (JsonProcessingException e) { + throw Throwables.propagate(e); + } + + String hashCode = DigestUtils.sha1Hex(dataSchema + + tuningConfig + + partitionOffsetStr + + minMsgTimeStr + + maxMsgTimeStr) + .substring(0, 15); + + return Joiner.on("_").join(baseTaskName(), dataSource, hashCode); + } + + protected abstract String baseTaskName(); + + private void updatePartitionDataFromStream() + { + Set partitionIds; + try { + synchronized (recordSupplierLock) { + partitionIds = recordSupplier.getPartitionIds(ioConfig.getStream()); + } + } + catch (Exception e) { + log.warn("Could not fetch partitions for topic/stream [%s]", ioConfig.getStream()); + log.debug(e, "full stack trace"); + return; + } + + if (partitionIds == null || partitionIds.size() == 0) { + log.warn("No partitions found for stream[%s]", ioConfig.getStream()); + return; + } + + log.debug("Found [%d] partitions for stream [%s]", partitionIds.size(), ioConfig.getStream()); + + Set closedPartitions = getOffsetsFromMetadataStorage() + .entrySet() + .stream() + .filter(x -> isEndOfShard(x.getValue())) + .map(Entry::getKey) + .collect(Collectors.toSet()); + + boolean initialPartitionDiscovery = this.partitionIds.isEmpty(); + for (PartitionIdType partitionId : partitionIds) { + if (closedPartitions.contains(partitionId)) { + log.info("partition [%s] is closed and has no more data, skipping.", partitionId); + continue; + } + + if (!initialPartitionDiscovery && !this.partitionIds.contains(partitionId)) { + subsequentlyDiscoveredPartitions.add(partitionId); + // should check for earlyPublishTime (Kinesis) here, not supported yet + } + + int taskGroupId = getTaskGroupIdForPartition(partitionId); + + ConcurrentHashMap partitionMap = partitionGroups.computeIfAbsent( + taskGroupId, + k -> new ConcurrentHashMap<>() + ); + + if (partitionMap.putIfAbsent(partitionId, getNotSetMarker()) == null) { + log.info( + "New partition [%s] discovered for stream [%s], added to task group [%d]", + partitionId, + ioConfig.getStream(), + taskGroupId + ); + } + } + } + + private void updateTaskStatus() throws ExecutionException, InterruptedException, TimeoutException + { + final List> futures = new ArrayList<>(); + final List futureTaskIds = new ArrayList<>(); + + // update status (and startTime if unknown) of current tasks in activelyReadingTaskGroups + for (TaskGroup group : activelyReadingTaskGroups.values()) { + for (Entry entry : group.tasks.entrySet()) { + final String taskId = entry.getKey(); + final TaskData taskData = entry.getValue(); + + if (taskData.startTime == null) { + futureTaskIds.add(taskId); + futures.add( + Futures.transform( + taskClient.getStartTimeAsync(taskId), new Function() + { + @Override + public Boolean apply(@Nullable DateTime startTime) + { + if (startTime == null) { + return false; + } + + taskData.startTime = startTime; + long millisRemaining = ioConfig.getTaskDuration().getMillis() - + (System.currentTimeMillis() - taskData.startTime.getMillis()); + if (millisRemaining > 0) { + scheduledExec.schedule( + buildRunTask(), + millisRemaining + MAX_RUN_FREQUENCY_MILLIS, + TimeUnit.MILLISECONDS + ); + } + + return true; + } + }, workerExec + ) + ); + } + + taskData.status = taskStorage.getStatus(taskId).get(); + } + } + + // update status of pending completion tasks in pendingCompletionTaskGroups + for (List taskGroups : pendingCompletionTaskGroups.values()) { + for (TaskGroup group : taskGroups) { + for (Entry entry : group.tasks.entrySet()) { + entry.getValue().status = taskStorage.getStatus(entry.getKey()).get(); + } + } + } + + List results = Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int i = 0; i < results.size(); i++) { + // false means the task hasn't started running yet and that's okay; null means it should be running but the HTTP + // request threw an exception so kill the task + if (results.get(i) == null) { + String taskId = futureTaskIds.get(i); + log.warn("Task [%s] failed to return start time, killing task", taskId); + killTask(taskId, "Task [%s] failed to return start time, killing task", taskId); + } + } + } + + private void checkTaskDuration() throws ExecutionException, InterruptedException, TimeoutException + { + final List>> futures = new ArrayList<>(); + final List futureGroupIds = new ArrayList<>(); + + for (Entry entry : activelyReadingTaskGroups.entrySet()) { + Integer groupId = entry.getKey(); + TaskGroup group = entry.getValue(); + + // find the longest running task from this group + DateTime earliestTaskStart = DateTimes.nowUtc(); + for (TaskData taskData : group.tasks.values()) { + if (taskData.startTime != null && earliestTaskStart.isAfter(taskData.startTime)) { + earliestTaskStart = taskData.startTime; + } + } + + // if this task has run longer than the configured duration, signal all tasks in the group to persist + if (earliestTaskStart.plus(ioConfig.getTaskDuration()).isBeforeNow()) { + log.info("Task group [%d] has run for [%s]", groupId, ioConfig.getTaskDuration()); + futureGroupIds.add(groupId); + futures.add(checkpointTaskGroup(group, true)); + } + } + + List> results = Futures.successfulAsList(futures) + .get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int j = 0; j < results.size(); j++) { + Integer groupId = futureGroupIds.get(j); + TaskGroup group = activelyReadingTaskGroups.get(groupId); + Map endOffsets = results.get(j); + + if (endOffsets != null) { + // set a timeout and put this group in pendingCompletionTaskGroups so that it can be monitored for completion + group.completionTimeout = DateTimes.nowUtc().plus(ioConfig.getCompletionTimeout()); + pendingCompletionTaskGroups.computeIfAbsent(groupId, k -> new CopyOnWriteArrayList<>()).add(group); + + // set endOffsets as the next startOffsets + for (Entry entry : endOffsets.entrySet()) { + partitionGroups.get(groupId).put(entry.getKey(), entry.getValue()); + } + } else { + for (String id : group.taskIds()) { + killTask( + id, + "All tasks in group [%s] failed to transition to publishing state", + groupId + ); + } + // clear partitionGroups, so that latest sequences from db is used as start sequences not the stale ones + // if tasks did some successful incremental handoffs + partitionGroups.get(groupId).replaceAll((partition, sequence) -> getNotSetMarker()); + } + + // remove this task group from the list of current task groups now that it has been handled + activelyReadingTaskGroups.remove(groupId); + } + } + + private ListenableFuture> checkpointTaskGroup( + final TaskGroup taskGroup, + final boolean finalize + ) + { + if (finalize) { + // 1) Check if any task completed (in which case we're done) and kill unassigned tasks + Iterator> i = taskGroup.tasks.entrySet().iterator(); + while (i.hasNext()) { + Entry taskEntry = i.next(); + String taskId = taskEntry.getKey(); + TaskData task = taskEntry.getValue(); + + if (task.status != null) { + if (task.status.isSuccess()) { + // If any task in this group has already completed, stop the rest of the tasks in the group and return. + // This will cause us to create a new set of tasks next cycle that will start from the sequences in + // metadata store (which will have advanced if we succeeded in publishing and will remain the same if + // publishing failed and we need to re-ingest) + return Futures.transform( + stopTasksInGroup(taskGroup, "task[%s] succeeded in the taskGroup", task.status.getId()), + new Function>() + { + @Nullable + @Override + public Map apply(@Nullable Object input) + { + return null; + } + } + ); + } + + if (task.status.isRunnable()) { + if (taskInfoProvider.getTaskLocation(taskId).equals(TaskLocation.unknown())) { + killTask(taskId, "Killing task [%s] which hasn't been assigned to a worker", taskId); + i.remove(); + } + } + } + } + } + + // 2) Pause running tasks + final List>> pauseFutures = new ArrayList<>(); + final List pauseTaskIds = ImmutableList.copyOf(taskGroup.taskIds()); + for (final String taskId : pauseTaskIds) { + pauseFutures.add(taskClient.pauseAsync(taskId)); + } + + return Futures.transform( + Futures.successfulAsList(pauseFutures), + new Function>, Map>() + { + @Nullable + @Override + public Map apply(List> input) + { + // 3) Build a map of the highest sequence read by any task in the group for each partition + final Map endOffsets = new HashMap<>(); + for (int i = 0; i < input.size(); i++) { + final Map result = input.get(i); + final String taskId = pauseTaskIds.get(i); + + if (result == null) { + // Get the exception + final Throwable pauseException; + try { + // The below get should throw ExecutionException since result is null. + final Map pauseResult = pauseFutures.get(i).get(); + throw new ISE( + "WTH? The pause request for task [%s] is supposed to fail, but returned [%s]", + taskId, + pauseResult + ); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } + catch (ExecutionException e) { + pauseException = e.getCause() == null ? e : e.getCause(); + } + + killTask( + taskId, + "An exception occured while waiting for task [%s] to pause: [%s]", + taskId, + pauseException + ); + taskGroup.tasks.remove(taskId); + + } else if (result.isEmpty()) { + killTask(taskId, "Task [%s] returned empty offsets after pause", taskId); + taskGroup.tasks.remove(taskId); + } else { // otherwise build a map of the highest sequences seen + for (Entry sequence : result.entrySet()) { + if (!endOffsets.containsKey(sequence.getKey()) + || makeSequenceNumber(endOffsets.get(sequence.getKey())).compareTo( + makeSequenceNumber(sequence.getValue())) < 0) { + endOffsets.put(sequence.getKey(), sequence.getValue()); + } + } + } + } + + // 4) Set the end sequences for each task to the values from step 3 and resume the tasks. All the tasks should + // finish reading and start publishing within a short period, depending on how in sync the tasks were. + final List> setEndOffsetFutures = new ArrayList<>(); + final List setEndOffsetTaskIds = ImmutableList.copyOf(taskGroup.taskIds()); + + if (setEndOffsetTaskIds.isEmpty()) { + log.info("All tasks in taskGroup [%d] have failed, tasks will be re-created", taskGroup.groupId); + return null; + } + + try { + + if (endOffsets.equals(taskGroup.checkpointSequences.lastEntry().getValue())) { + log.warn( + "Checkpoint [%s] is same as the start sequences [%s] of latest sequence for the task group [%d]", + endOffsets, + taskGroup.checkpointSequences.lastEntry().getValue(), + taskGroup.groupId + ); + } + + log.info( + "Setting endOffsets for tasks in taskGroup [%d] to %s and resuming", + taskGroup.groupId, + endOffsets + ); + for (final String taskId : setEndOffsetTaskIds) { + setEndOffsetFutures.add(taskClient.setEndOffsetsAsync(taskId, endOffsets, finalize)); + } + + List results = Futures.successfulAsList(setEndOffsetFutures) + .get(futureTimeoutInSeconds, TimeUnit.SECONDS); + for (int i = 0; i < results.size(); i++) { + if (results.get(i) == null || !results.get(i)) { + String taskId = setEndOffsetTaskIds.get(i); + killTask( + taskId, + "Task [%s] failed to respond to [set end offsets] in a timely manner, killing task", + taskId + ); + taskGroup.tasks.remove(taskId); + } + } + } + catch (Exception e) { + log.error("Something bad happened [%s]", e.getMessage()); + Throwables.propagate(e); + } + + if (taskGroup.tasks.isEmpty()) { + log.info("All tasks in taskGroup [%d] have failed, tasks will be re-created", taskGroup.groupId); + return null; + } + + return endOffsets; + } + }, + workerExec + ); + } + + private ListenableFuture stopTasksInGroup(@Nullable TaskGroup taskGroup, String stopReasonFormat, Object... args) + { + if (taskGroup == null) { + return Futures.immediateFuture(null); + } + + log.info( + "Stopping all tasks in taskGroup[%s] because: [%s]", + taskGroup.groupId, + StringUtils.format(stopReasonFormat, args) + ); + + final List> futures = new ArrayList<>(); + for (Entry entry : taskGroup.tasks.entrySet()) { + final String taskId = entry.getKey(); + final TaskData taskData = entry.getValue(); + if (taskData.status == null) { + killTask(taskId, "Killing task since task status is not known to supervisor"); + } else if (!taskData.status.isComplete()) { + futures.add(stopTask(taskId, false)); + } + } + + return Futures.successfulAsList(futures); + } + + private void checkPendingCompletionTasks() + throws ExecutionException, InterruptedException, TimeoutException + { + List> futures = new ArrayList<>(); + + for (Entry> pendingGroupList : pendingCompletionTaskGroups.entrySet()) { + + boolean stopTasksInTaskGroup = false; + Integer groupId = pendingGroupList.getKey(); + CopyOnWriteArrayList taskGroupList = pendingGroupList.getValue(); + List toRemove = new ArrayList<>(); + + for (TaskGroup group : taskGroupList) { + boolean foundSuccess = false, entireTaskGroupFailed = false; + + if (stopTasksInTaskGroup) { + // One of the earlier groups that was handling the same partition set timed out before the segments were + // published so stop any additional groups handling the same partition set that are pending completion. + futures.add( + stopTasksInGroup( + group, + "one of earlier groups that was handling the same partition set timed out before publishing segments" + ) + ); + toRemove.add(group); + continue; + } + + Iterator> iTask = group.tasks.entrySet().iterator(); + while (iTask.hasNext()) { + final Entry entry = iTask.next(); + final String taskId = entry.getKey(); + final TaskData taskData = entry.getValue(); + + Preconditions.checkNotNull(taskData.status, "WTH? task[%s] has a null status", taskId); + + if (taskData.status.isFailure()) { + iTask.remove(); // remove failed task + if (group.tasks.isEmpty()) { + // if all tasks in the group have failed, just nuke all task groups with this partition set and restart + entireTaskGroupFailed = true; + break; + } + } + + if (taskData.status.isSuccess()) { + // If one of the pending completion tasks was successful, stop the rest of the tasks in the group as + // we no longer need them to publish their segment. + log.info("Task [%s] completed successfully, stopping tasks %s", taskId, group.taskIds()); + futures.add( + stopTasksInGroup(group, "Task [%s] completed successfully, stopping tasks %s", taskId, group.taskIds()) + ); + foundSuccess = true; + toRemove.add(group); // remove the TaskGroup from the list of pending completion task groups + break; // skip iterating the rest of the tasks in this group as they've all been stopped now + } + } + + if ((!foundSuccess && group.completionTimeout.isBeforeNow()) || entireTaskGroupFailed) { + if (entireTaskGroupFailed) { + log.warn("All tasks in group [%d] failed to publish, killing all tasks for these partitions", groupId); + } else { + log.makeAlert( + "No task in [%s] for taskGroup [%d] succeeded before the completion timeout elapsed [%s]!", + group.taskIds(), + groupId, + ioConfig.getCompletionTimeout() + ).emit(); + } + + // reset partitions sequences for this task group so that they will be re-read from metadata storage + partitionGroups.get(groupId).replaceAll((partition, sequence) -> getNotSetMarker()); + // kill all the tasks in this pending completion group + killTasksInGroup( + group, + "No task in pending completion taskGroup[%d] succeeded before completion timeout elapsed", + groupId + ); + // set a flag so the other pending completion groups for this set of partitions will also stop + stopTasksInTaskGroup = true; + + // kill all the tasks in the currently reading task group and remove the bad task group + killTasksInGroup( + activelyReadingTaskGroups.remove(groupId), + "No task in the corresponding pending completion taskGroup[%d] succeeded before completion timeout elapsed", + groupId + ); + toRemove.add(group); + } + } + + taskGroupList.removeAll(toRemove); + } + + // wait for all task shutdowns to complete before returning + Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + + private void checkCurrentTaskState() throws ExecutionException, InterruptedException, TimeoutException + { + List> futures = new ArrayList<>(); + Iterator> iTaskGroups = activelyReadingTaskGroups.entrySet().iterator(); + while (iTaskGroups.hasNext()) { + Entry taskGroupEntry = iTaskGroups.next(); + Integer groupId = taskGroupEntry.getKey(); + TaskGroup taskGroup = taskGroupEntry.getValue(); + + // Iterate the list of known tasks in this group and: + // 1) Kill any tasks which are not "current" (have the partitions, starting sequences, and minimumMessageTime + // & maximumMessageTime (if applicable) in [activelyReadingTaskGroups]) + // 2) Remove any tasks that have failed from the list + // 3) If any task completed successfully, stop all the tasks in this group and move to the next group + + log.debug("Task group [%d] pre-pruning: %s", groupId, taskGroup.taskIds()); + + Iterator> iTasks = taskGroup.tasks.entrySet().iterator(); + while (iTasks.hasNext()) { + Entry task = iTasks.next(); + String taskId = task.getKey(); + TaskData taskData = task.getValue(); + + // stop and remove bad tasks from the task group + if (!isTaskCurrent(groupId, taskId)) { + log.info("Stopping task [%s] which does not match the expected sequence range and ingestion spec", taskId); + futures.add(stopTask(taskId, false)); + iTasks.remove(); + continue; + } + + Preconditions.checkNotNull(taskData.status, "WTH? task[%s] has a null status", taskId); + + // remove failed tasks + if (taskData.status.isFailure()) { + iTasks.remove(); + continue; + } + + // check for successful tasks, and if we find one, stop all tasks in the group and remove the group so it can + // be recreated with the next set of sequences + if (taskData.status.isSuccess()) { + futures.add(stopTasksInGroup(taskGroup, "task[%s] succeeded in the same taskGroup", taskData.status.getId())); + iTaskGroups.remove(); + break; + } + } + log.debug("Task group [%d] post-pruning: %s", groupId, taskGroup.taskIds()); + } + + // wait for all task shutdowns to complete before returning + Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + + private void createNewTasks() + throws JsonProcessingException + { + // update the checkpoints in the taskGroup to latest ones so that new tasks do not read what is already published + verifyAndMergeCheckpoints( + activelyReadingTaskGroups.values() + .stream() + .filter(taskGroup -> taskGroup.tasks.size() < ioConfig.getReplicas()) + .collect(Collectors.toList()) + ); + + // check that there is a current task group for each group of partitions in [partitionGroups] + for (Integer groupId : partitionGroups.keySet()) { + if (!activelyReadingTaskGroups.containsKey(groupId)) { + log.info("Creating new task group [%d] for partitions %s", groupId, partitionGroups.get(groupId).keySet()); + + Optional minimumMessageTime = (ioConfig.getLateMessageRejectionPeriod().isPresent() ? Optional.of( + DateTimes.nowUtc().minus(ioConfig.getLateMessageRejectionPeriod().get()) + ) : Optional.absent()); + + Optional maximumMessageTime = (ioConfig.getEarlyMessageRejectionPeriod().isPresent() ? Optional.of( + DateTimes.nowUtc().plus(ioConfig.getTaskDuration()).plus(ioConfig.getEarlyMessageRejectionPeriod().get()) + ) : Optional.absent()); + + + Map> startingOffsets = generateStartingSequencesForPartitionGroup( + groupId); + + ImmutableMap simpleStartingOffsets = startingOffsets + .entrySet() + .stream() + .filter(x -> x.getValue().get() != null) + .collect(Collectors.collectingAndThen( + Collectors.toMap(Entry::getKey, x -> x.getValue().get()), + ImmutableMap::copyOf + )); + + Set exclusiveStartSequenceNumberPartitions = !useExclusiveStartingSequence + ? Collections.emptySet() + : startingOffsets + .entrySet() + .stream() + .filter(x -> x.getValue().get() != null + && x.getValue().isExclusive()) + .map(Entry::getKey) + .collect(Collectors.toSet()); + + activelyReadingTaskGroups.put( + groupId, + new TaskGroup( + groupId, + simpleStartingOffsets, + minimumMessageTime, + maximumMessageTime, + exclusiveStartSequenceNumberPartitions + ) + ); + + } + } + + // iterate through all the current task groups and make sure each one has the desired number of replica tasks + boolean createdTask = false; + for (Entry entry : activelyReadingTaskGroups.entrySet()) { + TaskGroup taskGroup = entry.getValue(); + Integer groupId = entry.getKey(); + + if (taskGroup.startingSequences == null || taskGroup.startingSequences + .values().stream().allMatch(x -> x == null || isEndOfShard(x))) { + log.debug("Nothing to read in any partition for taskGroup [%d], skipping task creation", groupId); + continue; + } + + if (ioConfig.getReplicas() > taskGroup.tasks.size()) { + log.info( + "Number of tasks [%d] does not match configured numReplicas [%d] in task group [%d], creating more tasks", + taskGroup.tasks.size(), ioConfig.getReplicas(), groupId + ); + createTasksForGroup(groupId, ioConfig.getReplicas() - taskGroup.tasks.size()); + createdTask = true; + } + } + + if (createdTask && firstRunTime.isBeforeNow()) { + // Schedule a run event after a short delay to update our internal data structures with the new tasks that were + // just created. This is mainly for the benefit of the status API in situations where the run period is lengthy. + scheduledExec.schedule(buildRunTask(), 5000, TimeUnit.MILLISECONDS); + } + + } + + private void addNotice(Notice notice) + { + notices.add(notice); + } + + @VisibleForTesting + public void moveTaskGroupToPendingCompletion(int taskGroupId) + { + final TaskGroup taskGroup = activelyReadingTaskGroups.remove(taskGroupId); + if (taskGroup != null) { + pendingCompletionTaskGroups.computeIfAbsent(taskGroupId, k -> new CopyOnWriteArrayList<>()).add(taskGroup); + } + } + + @VisibleForTesting + public int getNoticesQueueSize() + { + return notices.size(); + } + + private ImmutableMap> generateStartingSequencesForPartitionGroup( + int groupId + ) + { + ImmutableMap.Builder> builder = ImmutableMap.builder(); + for (Entry entry : partitionGroups.get(groupId).entrySet()) { + PartitionIdType partition = entry.getKey(); + SequenceOffsetType sequence = entry.getValue(); + + if (!getNotSetMarker().equals(sequence)) { + // if we are given a startingOffset (set by a previous task group which is pending completion) then use it + if (!isEndOfShard(sequence)) { + builder.put(partition, makeSequenceNumber(sequence, false)); + } + } else { + // if we don't have a startingOffset (first run or we had some previous failures and reset the sequences) then + // get the sequence from metadata storage (if available) or Kafka/Kinesis (otherwise) + OrderedSequenceNumber offsetFromStorage = getOffsetFromStorageForPartition(partition); + + if (offsetFromStorage != null) { + builder.put(partition, offsetFromStorage); + } + } + } + return builder.build(); + } + + /** + * Queries the dataSource metadata table to see if there is a previous ending sequence for this partition. If it doesn't + * find any data, it will retrieve the latest or earliest Kafka/Kinesis sequence depending on the useEarliestOffset config. + */ + private OrderedSequenceNumber getOffsetFromStorageForPartition(PartitionIdType partition) + { + final Map metadataOffsets = getOffsetsFromMetadataStorage(); + SequenceOffsetType sequence = metadataOffsets.get(partition); + if (sequence != null) { + log.debug("Getting sequence [%s] from metadata storage for partition [%s]", sequence, partition); + if (!taskTuningConfig.isSkipSequenceNumberAvailabilityCheck()) { + if (!checkSequenceAvailability(partition, sequence)) { + if (taskTuningConfig.isResetOffsetAutomatically()) { + resetInternal( + createDataSourceMetaData(ioConfig.getStream(), ImmutableMap.of(partition, sequence)) + ); + throw new ISE( + "Previous sequenceNumber [%s] is no longer available for partition [%s] - automatically resetting sequence", + sequence, + partition + ); + + } else { + throw new ISE( + "Previous sequenceNumber [%s] is no longer available for partition [%s]. You can clear the previous sequenceNumber and start reading from a valid message by using the supervisor's reset API.", + sequence, + partition + ); + } + } + } + return makeSequenceNumber(sequence, true); + } else { + boolean useEarliestSequenceNumber = ioConfig.isUseEarliestSequenceNumber(); + if (subsequentlyDiscoveredPartitions.contains(partition)) { + log.info( + "Overriding useEarliestSequenceNumber and starting from beginning of newly discovered partition [%s] (which is probably from a split or merge)", + partition + ); + useEarliestSequenceNumber = true; + } + + sequence = getOffsetFromStreamForPartition(partition, useEarliestSequenceNumber); + if (sequence == null) { + throw new ISE("unable to fetch sequence number for partition[%s] from stream", partition); + } + log.info("Getting sequence number [%s] for partition [%s]", sequence, partition); + return makeSequenceNumber(sequence, false); + } + } + + private Map getOffsetsFromMetadataStorage() + { + final DataSourceMetadata dataSourceMetadata = indexerMetadataStorageCoordinator.getDataSourceMetadata(dataSource); + if (dataSourceMetadata instanceof SeekableStreamDataSourceMetadata + && checkSourceMetadataMatch(dataSourceMetadata)) { + @SuppressWarnings("unchecked") + SeekableStreamPartitions partitions = ((SeekableStreamDataSourceMetadata) dataSourceMetadata) + .getSeekableStreamPartitions(); + if (partitions != null) { + if (!ioConfig.getStream().equals(partitions.getStream())) { + log.warn( + "Topic/stream in metadata storage [%s] doesn't match spec topic/stream [%s], ignoring stored sequences", + partitions.getStream(), + ioConfig.getStream() + ); + return Collections.emptyMap(); + } else if (partitions.getPartitionSequenceNumberMap() != null) { + return partitions.getPartitionSequenceNumberMap(); + } + } + } + + return Collections.emptyMap(); + } + + @Nullable + private SequenceOffsetType getOffsetFromStreamForPartition(PartitionIdType partition, boolean useEarliestOffset) + { + synchronized (recordSupplierLock) { + StreamPartition topicPartition = new StreamPartition<>(ioConfig.getStream(), partition); + if (!recordSupplier.getAssignment().contains(topicPartition)) { + recordSupplier.assign(Collections.singleton(topicPartition)); + } + + return useEarliestOffset + ? recordSupplier.getEarliestSequenceNumber(topicPartition) + : recordSupplier.getLatestSequenceNumber(topicPartition); + } + } + + private void createTasksForGroup(int groupId, int replicas) + throws JsonProcessingException + { + TaskGroup group = activelyReadingTaskGroups.get(groupId); + Map startPartitions = group.startingSequences; + Map endPartitions = new HashMap<>(); + for (PartitionIdType partition : startPartitions.keySet()) { + endPartitions.put(partition, getEndOfPartitionMarker()); + } + Set exclusiveStartSequenceNumberPartitions = activelyReadingTaskGroups.get(groupId).exclusiveStartSequenceNumberPartitions; + + DateTime minimumMessageTime = activelyReadingTaskGroups.get(groupId).minimumMessageTime.orNull(); + DateTime maximumMessageTime = activelyReadingTaskGroups.get(groupId).maximumMessageTime.orNull(); + + SeekableStreamIndexTaskIOConfig newIoConfig = createTaskIoConfig( + groupId, + startPartitions, + endPartitions, + group.baseSequenceName, + minimumMessageTime, + maximumMessageTime, + exclusiveStartSequenceNumberPartitions, + ioConfig + ); + + + List> taskList = createIndexTasks( + replicas, + group.baseSequenceName, + sortingMapper, + group.checkpointSequences, + newIoConfig, + taskTuningConfig, + rowIngestionMetersFactory + ); + + for (SeekableStreamIndexTask indexTask : taskList) { + Optional taskQueue = taskMaster.getTaskQueue(); + if (taskQueue.isPresent()) { + try { + taskQueue.get().add(indexTask); + } + catch (EntryExistsException e) { + log.error("Tried to add task [%s] but it already exists", indexTask.getId()); + } + } else { + log.error("Failed to get task queue because I'm not the leader!"); + } + } + } + + @VisibleForTesting + public Runnable updateCurrentAndLatestOffsets() + { + return () -> { + try { + updateCurrentOffsets(); + updateLatestOffsetsFromStream(); + sequenceLastUpdated = DateTimes.nowUtc(); + } + catch (Exception e) { + log.warn(e, "Exception while getting current/latest sequences"); + } + }; + } + + private void updateCurrentOffsets() throws InterruptedException, ExecutionException, TimeoutException + { + final List> futures = Stream.concat( + activelyReadingTaskGroups.values().stream().flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()), + pendingCompletionTaskGroups.values() + .stream() + .flatMap(List::stream) + .flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()) + ).map( + task -> Futures.transform( + taskClient.getCurrentOffsetsAsync(task.getKey(), false), + (Function, Void>) (currentSequences) -> { + + if (currentSequences != null && !currentSequences.isEmpty()) { + task.getValue().currentSequences = currentSequences; + } + + return null; + } + ) + ).collect(Collectors.toList()); + + Futures.successfulAsList(futures).get(futureTimeoutInSeconds, TimeUnit.SECONDS); + } + + private void updateLatestOffsetsFromStream() throws InterruptedException + { + synchronized (recordSupplierLock) { + Set partitionIds = null; + try { + partitionIds = recordSupplier.getPartitionIds(ioConfig.getStream()); + } + catch (Exception e) { + log.warn("Could not fetch partitions for topic/stream [%s]", ioConfig.getStream()); + Throwables.propagate(e); + } + + Set> partitions = partitionIds + .stream() + .map(e -> new StreamPartition<>(ioConfig.getStream(), e)) + .collect(Collectors.toSet()); + + recordSupplier.assign(partitions); + recordSupplier.seekToLatest(partitions); + + updateLatestSequenceFromStream(recordSupplier, partitions); + } + + } + + protected abstract void updateLatestSequenceFromStream( + RecordSupplier recordSupplier, + Set> partitions + ); + + protected Map getHighestCurrentOffsets() + { + return activelyReadingTaskGroups + .values() + .stream() + .flatMap(taskGroup -> taskGroup.tasks.entrySet().stream()) + .flatMap(taskData -> taskData.getValue().currentSequences.entrySet().stream()) + .collect(Collectors.toMap( + Entry::getKey, + Entry::getValue, + (v1, v2) -> makeSequenceNumber(v1).compareTo(makeSequenceNumber(v2)) > 0 ? v1 : v2 + )); + } + + private OrderedSequenceNumber makeSequenceNumber(SequenceOffsetType seq) + { + return makeSequenceNumber(seq, false); + } + + + // exposed for testing for visibility into initialization state + @VisibleForTesting + public boolean isStarted() + { + return started; + } + + // exposed for testing for visibility into initialization state + @VisibleForTesting + public boolean isLifecycleStarted() + { + return lifecycleStarted; + } + + // exposed for testing for visibility into initialization state + @VisibleForTesting + public int getInitRetryCounter() + { + return initRetryCounter; + } + + // exposed for testing to allow "bootstrap.servers" to be changed after supervisor is created + @VisibleForTesting + public SeekableStreamSupervisorIOConfig getIoConfig() + { + return ioConfig; + } + + @Override + public void checkpoint( + @Nullable Integer taskGroupId, + @Deprecated String baseSequenceName, + DataSourceMetadata previousCheckPoint, + DataSourceMetadata currentCheckPoint + ) + { + Preconditions.checkNotNull(previousCheckPoint, "previousCheckpoint"); + Preconditions.checkNotNull(currentCheckPoint, "current checkpoint cannot be null"); + Preconditions.checkArgument( + spec.getIoConfig() + .getStream() + .equals(((SeekableStreamDataSourceMetadata) currentCheckPoint).getSeekableStreamPartitions() + .getStream()), + "Supervisor stream [%s] and stream in checkpoint [%s] does not match", + spec.getIoConfig().getStream(), + ((SeekableStreamDataSourceMetadata) currentCheckPoint).getSeekableStreamPartitions().getStream() + ); + + log.info("Checkpointing [%s] for taskGroup [%s]", currentCheckPoint, taskGroupId); + addNotice( + new CheckpointNotice( + taskGroupId, + baseSequenceName, + (SeekableStreamDataSourceMetadata) previousCheckPoint, + (SeekableStreamDataSourceMetadata) currentCheckPoint + ) + ); + } + + /** + * creates a specific task IOConfig instance for Kafka/Kinesis + * + * @return specific instance of Kafka/Kinesis IOConfig + */ + protected abstract SeekableStreamIndexTaskIOConfig createTaskIoConfig( + int groupId, + Map startPartitions, + Map endPartitions, + String baseSequenceName, + DateTime minimumMessageTime, + DateTime maximumMessageTime, + Set exclusiveStartSequenceNumberPartitions, + SeekableStreamSupervisorIOConfig ioConfig + ); + + /** + * creates a list of specific kafka/kinesis index tasks using + * the given replicas count + * + * @return list of specific kafka/kinesis index taksks + * + * @throws JsonProcessingException + */ + protected abstract List> createIndexTasks( + int replicas, + String baseSequenceName, + ObjectMapper sortingMapper, + TreeMap> sequenceOffsets, + SeekableStreamIndexTaskIOConfig taskIoConfig, + SeekableStreamIndexTaskTuningConfig taskTuningConfig, + RowIngestionMetersFactory rowIngestionMetersFactory + ) throws JsonProcessingException; + + /** + * calculates the taskgroup id that the given partition belongs to. + * different between Kafka/Kinesis since Kinesis uses String as partition id + * + * @param partition paritition id + * + * @return taskgroup id + */ + protected abstract int getTaskGroupIdForPartition(PartitionIdType partition); + + /** + * checks if the passed in DataSourceMetadata is a specific instance + * of [kafka/kinesis]DataSourceMetadata + * + * @param metadata datasource metadata + * + * @return true if isInstance else false + */ + protected abstract boolean checkSourceMetadataMatch(DataSourceMetadata metadata); + + /** + * checks if the passed in Task is a specific instance of + * [Kafka/Kinesis]IndexTask + * + * @param task task + * + * @return true if isInstance else false + */ + protected abstract boolean doesTaskTypeMatchSupervisor(Task task); + + /** + * creates a specific instance of kafka/kinesis datasource metadata + * + * @param stream stream name + * @param map partitionId -> sequence + * + * @return specific instance of datasource metadata + */ + protected abstract SeekableStreamDataSourceMetadata createDataSourceMetaData( + String stream, + Map map + ); + + /** + * wraps the passed in SequenceOffsetType sequence number into a {@link OrderedSequenceNumber} object + * to facilitate comparison and accomodate exclusive starting sequennce in kinesis + * + * @return specific instance of [Kafka/Kinesis]OrderedSequenceNumber + */ + protected abstract OrderedSequenceNumber makeSequenceNumber( + SequenceOffsetType seq, + boolean isExclusive + ); + + /** + * schedules periodic emitLag() reporting for Kafka, not yet implemented in Kinesis, + * but will be in the future + */ + protected abstract void scheduleReporting(ScheduledExecutorService reportingExec); + + /** + * calculate lag per partition for kafka, kinesis implementation returns an empty + * map + * + * @return map of partition id -> lag + */ + protected abstract Map getLagPerPartition(Map currentOffsets); + + /** + * returns an instance of a specific Kinesis/Kafka recordSupplier + * + * @return specific instance of Kafka/Kinesis RecordSupplier + */ + protected abstract RecordSupplier setupRecordSupplier(); + + /** + * creates a specific instance of Kafka/Kinesis Supervisor Report Payload + * + * @return specific instance of Kafka/Kinesis Supervisor Report Payload + */ + protected abstract SeekableStreamSupervisorReportPayload createReportPayload( + int numPartitions, + boolean includeOffsets + ); + + /** + * checks if sequence from metadata storage is still valid + * + * @return true if still valid else false + */ + private boolean checkSequenceAvailability( + @NotNull PartitionIdType partition, + @NotNull SequenceOffsetType sequenceFromMetadata + ) + { + SequenceOffsetType earliestSequence = getOffsetFromStreamForPartition(partition, true); + SequenceOffsetType latestSequence = getOffsetFromStreamForPartition(partition, false); + return (earliestSequence == null + || makeSequenceNumber(earliestSequence).compareTo(makeSequenceNumber(sequenceFromMetadata)) <= 0) + && (latestSequence == null + || makeSequenceNumber(latestSequence).compareTo(makeSequenceNumber(sequenceFromMetadata)) >= 0); + } + + + /** + * a special sequence number that is used to indicate that the sequence offset + * for a particular partition has not yet been calculated by the supervisor. When + * the not_set marker is read by the supervisor, it will first attempt to restore it + * from metadata storage, if that fails, from the Kafka/Kinesis + * + * @return sequence offset that represets NOT_SET + */ + protected abstract SequenceOffsetType getNotSetMarker(); + + /** + * returns the logical maximum number for a Kafka partition or Kinesis shard. This is + * used to set the initial endoffsets when creating a new task, since we don't know + * what sequence offsets to read to initially + * + * @return end of partition sequence offset + */ + protected abstract SequenceOffsetType getEndOfPartitionMarker(); + + /** + * checks if seqNum marks the end of a Kinesis shard. Used by Kinesis only. + */ + protected abstract boolean isEndOfShard(SequenceOffsetType seqNum); +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorIOConfig.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorIOConfig.java new file mode 100644 index 000000000000..e62f675d7fc5 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorIOConfig.java @@ -0,0 +1,134 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.supervisor; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Optional; +import com.google.common.base.Preconditions; +import org.joda.time.Duration; +import org.joda.time.Period; + +public abstract class SeekableStreamSupervisorIOConfig +{ + private final String stream; + private final Integer replicas; + private final Integer taskCount; + private final Duration taskDuration; + private final Duration startDelay; + private final Duration period; + private final boolean useEarliestSequenceNumber; + private final Duration completionTimeout; + private final Optional lateMessageRejectionPeriod; + private final Optional earlyMessageRejectionPeriod; + + public SeekableStreamSupervisorIOConfig( + String stream, + Integer replicas, + Integer taskCount, + Period taskDuration, + Period startDelay, + Period period, + Boolean useEarliestSequenceNumber, + Period completionTimeout, + Period lateMessageRejectionPeriod, + Period earlyMessageRejectionPeriod + ) + { + this.stream = Preconditions.checkNotNull(stream, "stream cannot be null"); + this.replicas = replicas != null ? replicas : 1; + this.taskCount = taskCount != null ? taskCount : 1; + this.taskDuration = defaultDuration(taskDuration, "PT1H"); + this.startDelay = defaultDuration(startDelay, "PT5S"); + this.period = defaultDuration(period, "PT30S"); + this.useEarliestSequenceNumber = useEarliestSequenceNumber != null ? useEarliestSequenceNumber : false; + this.completionTimeout = defaultDuration(completionTimeout, "PT30M"); + this.lateMessageRejectionPeriod = lateMessageRejectionPeriod == null + ? Optional.absent() + : Optional.of(lateMessageRejectionPeriod.toStandardDuration()); + this.earlyMessageRejectionPeriod = earlyMessageRejectionPeriod == null + ? Optional.absent() + : Optional.of(earlyMessageRejectionPeriod.toStandardDuration()); + } + + private static Duration defaultDuration(final Period period, final String theDefault) + { + return (period == null ? new Period(theDefault) : period).toStandardDuration(); + } + + @JsonProperty + public String getStream() + { + return stream; + } + + @JsonProperty + public Integer getReplicas() + { + return replicas; + } + + @JsonProperty + public Integer getTaskCount() + { + return taskCount; + } + + @JsonProperty + public Duration getTaskDuration() + { + return taskDuration; + } + + @JsonProperty + public Duration getStartDelay() + { + return startDelay; + } + + @JsonProperty + public Duration getPeriod() + { + return period; + } + + @JsonProperty + public boolean isUseEarliestSequenceNumber() + { + return useEarliestSequenceNumber; + } + + @JsonProperty + public Duration getCompletionTimeout() + { + return completionTimeout; + } + + @JsonProperty + public Optional getEarlyMessageRejectionPeriod() + { + return earlyMessageRejectionPeriod; + } + + @JsonProperty + public Optional getLateMessageRejectionPeriod() + { + return lateMessageRejectionPeriod; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorReportPayload.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorReportPayload.java new file mode 100644 index 000000000000..29547c335a66 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorReportPayload.java @@ -0,0 +1,157 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.supervisor; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.java.util.common.IAE; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +@JsonInclude(JsonInclude.Include.NON_NULL) +public abstract class SeekableStreamSupervisorReportPayload +{ + private final String dataSource; + private final String stream; + private final int partitions; + private final int replicas; + private final long durationSeconds; + private final List activeTasks; + private final List publishingTasks; + private final Map latestOffsets; + private final Map minimumLag; + private final Long aggregateLag; + private final DateTime offsetsLastUpdated; + private final boolean suspended; + + public SeekableStreamSupervisorReportPayload( + String dataSource, + String stream, + int partitions, + int replicas, + long durationSeconds, + @Nullable Map latestOffsets, + @Nullable Map minimumLag, + @Nullable Long aggregateLag, + @Nullable DateTime offsetsLastUpdated, + boolean suspended + ) + { + this.dataSource = dataSource; + this.stream = stream; + this.partitions = partitions; + this.replicas = replicas; + this.durationSeconds = durationSeconds; + this.activeTasks = new ArrayList<>(); + this.publishingTasks = new ArrayList<>(); + this.latestOffsets = latestOffsets; + this.minimumLag = minimumLag; + this.aggregateLag = aggregateLag; + this.offsetsLastUpdated = offsetsLastUpdated; + this.suspended = suspended; + } + + public void addTask(TaskReportData data) + { + if (data.getType().equals(TaskReportData.TaskType.ACTIVE)) { + activeTasks.add(data); + } else if (data.getType().equals(TaskReportData.TaskType.PUBLISHING)) { + publishingTasks.add(data); + } else { + throw new IAE("Unknown task type [%s]", data.getType().name()); + } + } + + @JsonProperty + public String getDataSource() + { + return dataSource; + } + + @JsonProperty + public String getStream() + { + return stream; + } + + @JsonProperty + public int getPartitions() + { + return partitions; + } + + @JsonProperty + public int getReplicas() + { + return replicas; + } + + @JsonProperty + public boolean getSuspended() + { + return suspended; + } + + @JsonProperty + public long getDurationSeconds() + { + return durationSeconds; + } + + @JsonProperty + public List getActiveTasks() + { + return activeTasks; + } + + @JsonProperty + public List getPublishingTasks() + { + return publishingTasks; + } + + @JsonProperty + public Map getLatestOffsets() + { + return latestOffsets; + } + + @JsonProperty + public Map getMinimumLag() + { + return minimumLag; + } + + @JsonProperty + public Long getAggregateLag() + { + return aggregateLag; + } + + @JsonProperty + public DateTime getOffsetsLastUpdated() + { + return offsetsLastUpdated; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorSpec.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorSpec.java new file mode 100644 index 000000000000..961e066c94a3 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorSpec.java @@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.supervisor; + +import com.fasterxml.jackson.annotation.JacksonInject; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.indexing.common.stats.RowIngestionMetersFactory; +import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.supervisor.Supervisor; +import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.server.metrics.DruidMonitorSchedulerConfig; + +import java.util.List; +import java.util.Map; + +public abstract class SeekableStreamSupervisorSpec implements SupervisorSpec +{ + protected final TaskStorage taskStorage; + protected final TaskMaster taskMaster; + protected final IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator; + protected final SeekableStreamIndexTaskClientFactory indexTaskClientFactory; + protected final ObjectMapper mapper; + protected final RowIngestionMetersFactory rowIngestionMetersFactory; + private final DataSchema dataSchema; + private final SeekableStreamSupervisorTuningConfig tuningConfig; + private final SeekableStreamSupervisorIOConfig ioConfig; + private final Map context; + protected final ServiceEmitter emitter; + protected final DruidMonitorSchedulerConfig monitorSchedulerConfig; + private final boolean suspended; + + @JsonCreator + public SeekableStreamSupervisorSpec( + @JsonProperty("dataSchema") DataSchema dataSchema, + @JsonProperty("tuningConfig") SeekableStreamSupervisorTuningConfig tuningConfig, + @JsonProperty("ioConfig") SeekableStreamSupervisorIOConfig ioConfig, + @JsonProperty("context") Map context, + @JsonProperty("suspended") Boolean suspended, + @JacksonInject TaskStorage taskStorage, + @JacksonInject TaskMaster taskMaster, + @JacksonInject IndexerMetadataStorageCoordinator indexerMetadataStorageCoordinator, + @JacksonInject SeekableStreamIndexTaskClientFactory indexTaskClientFactory, + @JacksonInject @Json ObjectMapper mapper, + @JacksonInject ServiceEmitter emitter, + @JacksonInject DruidMonitorSchedulerConfig monitorSchedulerConfig, + @JacksonInject RowIngestionMetersFactory rowIngestionMetersFactory + ) + { + this.dataSchema = Preconditions.checkNotNull(dataSchema, "dataSchema"); + this.tuningConfig = tuningConfig; // null check done in concrete class + this.ioConfig = Preconditions.checkNotNull(ioConfig, "ioConfig"); + this.context = context; + + this.taskStorage = taskStorage; + this.taskMaster = taskMaster; + this.indexerMetadataStorageCoordinator = indexerMetadataStorageCoordinator; + this.indexTaskClientFactory = indexTaskClientFactory; + this.mapper = mapper; + this.emitter = emitter; + this.monitorSchedulerConfig = monitorSchedulerConfig; + this.rowIngestionMetersFactory = rowIngestionMetersFactory; + this.suspended = suspended != null ? suspended : false; + } + + @JsonProperty + public DataSchema getDataSchema() + { + return dataSchema; + } + + @JsonProperty + public SeekableStreamSupervisorTuningConfig getTuningConfig() + { + return tuningConfig; + } + + @JsonProperty + public SeekableStreamSupervisorIOConfig getIoConfig() + { + return ioConfig; + } + + @JsonProperty + public Map getContext() + { + return context; + } + + public ServiceEmitter getEmitter() + { + return emitter; + } + + @Override + public String getId() + { + return dataSchema.getDataSource(); + } + + public DruidMonitorSchedulerConfig getMonitorSchedulerConfig() + { + return monitorSchedulerConfig; + } + + @Override + public abstract Supervisor createSupervisor(); + + @Override + public List getDataSources() + { + return ImmutableList.of(getDataSchema().getDataSource()); + } + + @Override + public SeekableStreamSupervisorSpec createSuspendedSpec() + { + return toggleSuspend(true); + } + + @Override + public SeekableStreamSupervisorSpec createRunningSpec() + { + return toggleSuspend(false); + } + + @Override + @JsonProperty("suspended") + public boolean isSuspended() + { + return suspended; + } + + protected abstract SeekableStreamSupervisorSpec toggleSuspend(boolean suspend); + + +} diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorTuningConfig.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorTuningConfig.java new file mode 100644 index 000000000000..528bfb5efe90 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorTuningConfig.java @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.supervisor; + +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig; +import org.joda.time.Duration; +import org.joda.time.Period; + +public interface SeekableStreamSupervisorTuningConfig +{ + + int DEFAULT_CHAT_RETRIES = 8; + String DEFAULT_HTTP_TIMEOUT = "PT10S"; + String DEFAULT_SHUTDOWN_TIMEOUT = "PT80S"; + + + static Duration defaultDuration(final Period period, final String theDefault) + { + return (period == null ? new Period(theDefault) : period).toStandardDuration(); + } + + @JsonProperty + Integer getWorkerThreads(); + + @JsonProperty + Integer getChatThreads(); + + @JsonProperty + Long getChatRetries(); + + @JsonProperty + Duration getHttpTimeout(); + + @JsonProperty + Duration getShutdownTimeout(); + + SeekableStreamIndexTaskTuningConfig convertToTaskTuningConfig(); +} diff --git a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/TaskReportData.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/TaskReportData.java similarity index 72% rename from extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/TaskReportData.java rename to indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/TaskReportData.java index 85b05c7a5f0d..0b18e58ee678 100644 --- a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/TaskReportData.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/TaskReportData.java @@ -17,7 +17,7 @@ * under the License. */ -package org.apache.druid.indexing.kafka.supervisor; +package org.apache.druid.indexing.seekablestream.supervisor; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; @@ -26,29 +26,24 @@ import javax.annotation.Nullable; import java.util.Map; -public class TaskReportData +public class TaskReportData { - public enum TaskType - { - ACTIVE, PUBLISHING, UNKNOWN - } - private final String id; - private final Map startingOffsets; + private final Map startingOffsets; private final DateTime startTime; private final Long remainingSeconds; private final TaskType type; - private final Map currentOffsets; - private final Map lag; + private final Map currentOffsets; + private final Map lag; public TaskReportData( String id, - @Nullable Map startingOffsets, - @Nullable Map currentOffsets, + @Nullable Map startingOffsets, + @Nullable Map currentOffsets, @Nullable DateTime startTime, Long remainingSeconds, TaskType type, - @Nullable Map lag + @Nullable Map lag ) { this.id = id; @@ -68,14 +63,14 @@ public String getId() @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) - public Map getStartingOffsets() + public Map getStartingOffsets() { return startingOffsets; } @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) - public Map getCurrentOffsets() + public Map getCurrentOffsets() { return currentOffsets; } @@ -100,7 +95,7 @@ public TaskType getType() @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) - public Map getLag() + public Map getLag() { return lag; } @@ -117,4 +112,15 @@ public String toString() (lag != null ? ", lag=" + lag : "") + '}'; } + + /** + * Used by the Supervisor to report status of tasks + * ACTIVE - task is waiting to be started, started, or reading + * PUBLISHING - task is publishing or registering handoff + * UNNKNOWN - unknown + */ + public enum TaskType + { + ACTIVE, PUBLISHING, UNKNOWN + } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/utils/RandomIdUtils.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/utils/RandomIdUtils.java new file mode 100644 index 000000000000..4715db8d3f35 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/utils/RandomIdUtils.java @@ -0,0 +1,37 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.seekablestream.utils; + +import java.util.Random; +import java.util.concurrent.ThreadLocalRandom; + +public class RandomIdUtils +{ + private static final Random RANDOM = ThreadLocalRandom.current(); + + public static String getRandomId() + { + final StringBuilder suffix = new StringBuilder(8); + for (int i = 0; i < Integer.BYTES * 2; ++i) { + suffix.append((char) ('a' + ((RANDOM.nextInt() >>> (i * 4)) & 0x0F))); + } + return suffix.toString(); + } +} diff --git a/pom.xml b/pom.xml index ca2321486f5a..e861985419d1 100644 --- a/pom.xml +++ b/pom.xml @@ -134,6 +134,7 @@ extensions-core/kafka-eight extensions-core/kafka-extraction-namespace extensions-core/kafka-indexing-service + extensions-core/kinesis-indexing-service extensions-core/mysql-metadata-storage extensions-core/parquet-extensions extensions-core/postgresql-metadata-storage @@ -950,7 +951,9 @@ ${project.parent.basedir}/codestyle/joda-time-forbidden-apis.txt ${project.parent.basedir}/codestyle/druid-forbidden-apis.txt - **.SuppressForbidden + + **.SuppressForbidden + @@ -1000,6 +1003,8 @@ sun.nio.ch.DirectBuffer sun.misc.Cleaner sun.misc.Unsafe + + java.lang.invoke.MethodHandle @@ -1349,9 +1354,9 @@ -Xmx768m -Duser.language=en -Duser.country=US -Dfile.encoding=UTF-8 - -Duser.timezone=UTC -Djava.util.logging.manager=org.apache.logging.log4j.jul.LogManager - - -Ddruid.indexing.doubleStorage=double + -Duser.timezone=UTC -Djava.util.logging.manager=org.apache.logging.log4j.jul.LogManager + + -Ddruid.indexing.doubleStorage=double true