/** * Produces a bucket as a record to appropriate bulk action stream. */ protected void produceBucket(ComputationContext context, String action, String commandId, int bucketSize, long bucketNumber) { List<String> ids = documentIds.subList(0, min(bucketSize, documentIds.size())); BulkBucket bucket = new BulkBucket(commandId, ids); String key = commandId + ":" + Long.toString(bucketNumber); Record record = Record.of(key, BulkCodecs.getBucketCodec().encode(bucket)); if (produceImmediate) { ((ComputationContextImpl) context).produceRecordImmediate(action, record); } else { context.produceRecord(action, record); } ids.clear(); // this clear the documentIds part that has been sent }
public static void updateStatus(ComputationContext context, BulkStatus delta) { context.produceRecord(OUTPUT_1, delta.getId(), BulkCodecs.getStatusCodec().encode(delta)); }
public LogOffset append(int partition, String key, M message) { Bytes value = Bytes.wrap(encodingCodec.encode(message)); ProducerRecord<String, Bytes> record = new ProducerRecord<>(topic, partition, key, value); Future<RecordMetadata> future = producer.send(record); RecordMetadata result; try { result = future.get(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new StreamRuntimeException("Unable to send record: " + record, e); } catch (ExecutionException e) { throw new StreamRuntimeException("Unable to send record: " + record, e); } LogOffset ret = new LogOffsetImpl(name, partition, result.offset()); if (log.isDebugEnabled()) { int len = record.value().get().length; log.debug(String.format("append to %s-%02d:+%d, len: %d, key: %s, value: %s", name, partition, ret.offset(), len, key, message)); } return ret; }
protected void updateStatusAsScrolling(ComputationContext context, String commandId) { BulkStatus delta = BulkStatus.deltaOf(commandId); delta.setState(SCROLLING_RUNNING); delta.setScrollStartTime(Instant.now()); ((ComputationContextImpl) context).produceRecordImmediate(STATUS_STREAM, commandId, BulkCodecs.getStatusCodec().encode(delta)); }
/** * Stores the command in the kv store, returns the encoded command. */ public byte[] setCommand(BulkCommand command) { KeyValueStore kvStore = getKvStore(); byte[] commandAsBytes = BulkCodecs.getCommandCodec().encode(command); kvStore.put(COMMAND_PREFIX + command.getId(), commandAsBytes); return commandAsBytes; }
@Override public LogOffset append(int partition, M message) { ExcerptAppender appender = partitions.get(partition).acquireAppender(); if (NO_CODEC.equals(codec)) { // default format for backward compatibility appender.writeDocument(w -> w.write(MSG_KEY).object(message)); } else { appender.writeDocument(w -> w.write().bytes(codec.encode(message))); } long offset = appender.lastIndexAppended(); LogOffset ret = new LogOffsetImpl(name, partition, offset); if (log.isDebugEnabled()) { log.debug(String.format("append to %s, value: %s", ret, message)); } return ret; }
protected void finishBlob(ComputationContext context, String commandId) { String outputStream = getOutputStream(commandId); DataBucket in = lastBuckets.get(commandId); if (!SORT_STREAM.equals(outputStream)) { appendHeaderFooterToFile(createTemp(commandId), commandId, in.getHeader(), in.getFooter()); } String storeName = Framework.getService(BulkService.class).getStatus(commandId).getAction(); String value = saveInTransientStore(commandId, storeName); DataBucket out = new DataBucket(commandId, totals.get(commandId), value, in.getHeaderAsString(), in.getFooterAsString()); Codec<DataBucket> codec = BulkCodecs.getDataBucketCodec(); if (produceImmediate) { ((ComputationContextImpl) context).produceRecordImmediate(outputStream, Record.of(commandId, codec.encode(out))); } else { context.produceRecord(outputStream, Record.of(commandId, codec.encode(out))); } totals.remove(commandId); counters.remove(commandId); lastBuckets.remove(commandId); // we checkpoint only if there is not another command in progress if (counters.isEmpty()) { context.askForCheckpoint(); } }
protected void updateStatusAfterScroll(ComputationContext context, String commandId, long documentCount, String errorMessage) { BulkStatus delta = BulkStatus.deltaOf(commandId); if (errorMessage != null) { delta.inError(errorMessage); } if (documentCount == 0) { delta.setState(COMPLETED); delta.setCompletedTime(Instant.now()); } else { delta.setState(RUNNING); } delta.setScrollEndTime(Instant.now()); delta.setTotal(documentCount); ((ComputationContextImpl) context).produceRecordImmediate(STATUS_STREAM, commandId, BulkCodecs.getStatusCodec().encode(delta)); }
@Override public void endBucket(ComputationContext context, BulkStatus delta) { long bucketSize = delta.getProcessed(); bulkRequests.add(bulkRequest); String commandId = getCurrentCommand().getId(); int i = 0; int count = 0; for (BulkRequest request : bulkRequests) { DataBucket dataBucket = new DataBucket(commandId, request.numberOfActions(), toBytes(request)); // use distinct key to distribute the message evenly between partitions String key = bucketKey + "-" + i++; context.produceRecord(OUTPUT_1, Record.of(key, BulkCodecs.getDataBucketCodec().encode(dataBucket))); count += request.numberOfActions(); } if (count < bucketSize) { log.warn(String.format("Command: %s offset: %s created %d documents out of %d, %d not accessible", commandId, context.getLastOffset(), count, bucketSize, bucketSize - count)); DataBucket dataBucket = new DataBucket(commandId, bucketSize - count, toBytes(new BulkRequest())); context.produceRecord(OUTPUT_1, Record.of(bucketKey + "-missing", BulkCodecs.getDataBucketCodec().encode(dataBucket))); } bulkRequest = null; bulkRequests.clear(); }
context.produceRecord(outputStream, Record.of(commandId, codec.encode(out))); context.askForCheckpoint();
@Override public BulkStatus abort(String commandId) { BulkStatus status = getStatus(commandId); if (COMPLETED.equals(status.getState())) { log.debug("Cannot abort a completed command: {}", commandId); return status; } status.setState(ABORTED); // set the status in the KV store setStatus(status); // Send a delta to the status computation BulkStatus delta = BulkStatus.deltaOf(commandId); delta.setCompletedTime(Instant.now()); delta.setState(ABORTED); byte[] statusAsBytes = BulkCodecs.getStatusCodec().encode(delta); LogManager logManager = Framework.getService(StreamService.class).getLogManager(BULK_LOG_MANAGER_NAME); LogAppender<Record> logAppender = logManager.getAppender(STATUS_STREAM); logAppender.append(commandId, Record.of(commandId, statusAsBytes)); return status; }
/** * Stores the status in the kv store returns the encoded status */ public byte[] setStatus(BulkStatus status) { KeyValueStore kvStore = getKvStore(); byte[] statusAsBytes = BulkCodecs.getStatusCodec().encode(status); switch (status.getState()) { case ABORTED: kvStore.put(STATUS_PREFIX + status.getId(), statusAsBytes, ABORTED_TTL_SECONDS); // we remove the command from the kv store, so computation have to handle abort kvStore.put(COMMAND_PREFIX + status.getId(), (String) null); break; case COMPLETED: kvStore.put(STATUS_PREFIX + status.getId(), statusAsBytes, COMPLETED_TTL_SECONDS); kvStore.setTTL(COMMAND_PREFIX + status.getId(), COMPLETED_TTL_SECONDS); break; default: kvStore.put(STATUS_PREFIX + status.getId(), statusAsBytes); } return statusAsBytes; }
@Override public void processRecord(ComputationContext context, String inputStreamName, Record record) { Codec<DataBucket> codec = BulkCodecs.getDataBucketCodec(); DataBucket in = codec.decode(record.getData()); String storeName = Framework.getService(BulkService.class).getStatus(in.getCommandId()).getAction(); Blob blob = getBlob(in.getDataAsString(), storeName); try { blob = BlobUtils.zip(blob, blob.getFilename() + ".zip"); } catch (IOException e) { log.error("Unable to zip blob", e); } storeBlob(blob, in.getCommandId(), storeName); DataBucket out = new DataBucket(in.getCommandId(), in.getCount(), getTransientStoreKey(in.getCommandId())); context.produceRecord(OUTPUT_1, Record.of(in.getCommandId(), codec.encode(out))); context.askForCheckpoint(); }