@Override public BulkCommand getCommand(String commandId) { KeyValueStore keyValueStore = getKvStore(); byte[] statusAsBytes = keyValueStore.get(COMMAND_PREFIX + commandId); if (statusAsBytes == null) { return null; } return BulkCodecs.getCommandCodec().decode(statusAsBytes); }
/** * 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 }
@SuppressWarnings("unchecked") protected Codec<Record> getCodecForStreams(String name, Set<String> streams) { Codec<Record> codec = null; Set<String> codecNames = new HashSet<>(); for (String stream : streams) { codec = settings.getCodec(stream); codecNames.add(codec == null ? "none" : codec.getName()); } if (codecNames.size() > 1) { throw new IllegalArgumentException(String.format("Different codecs for computation %s: %s", name, Arrays.toString(codecNames.toArray()))); } if (codec == null) { codec = NO_CODEC; } return codec; }
@Override public void processRecord(ComputationContext context, String inputStreamName, Record record) { Codec<DataBucket> codec = BulkCodecs.getDataBucketCodec(); DataBucket in = codec.decode(record.getData()); context.produceRecord(outputStream, Record.of(commandId, codec.encode(out))); context.askForCheckpoint();
@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(); }
@Override public BulkStatus getStatus(String commandId) { KeyValueStore keyValueStore = getKvStore(); byte[] statusAsBytes = keyValueStore.get(STATUS_PREFIX + commandId); if (statusAsBytes == null) { log.debug("Request status of unknown command: {}", commandId); return BulkStatus.unknownOf(commandId); } return BulkCodecs.getStatusCodec().decode(statusAsBytes); }
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; }
@Override public void processRecord(ComputationContext context, String documentIdsStreamName, Record record) { Codec<DataBucket> codec = BulkCodecs.getDataBucketCodec(); DataBucket in = codec.decode(record.getData()); String commandId = in.getCommandId(); long nbDocuments = in.getCount(); appendToFile(commandId, in.getData()); if (counters.containsKey(commandId)) { counters.put(commandId, nbDocuments + counters.get(commandId)); } else { counters.put(commandId, nbDocuments); } lastBuckets.put(commandId, in); if (counters.get(commandId) < getTotal(commandId)) { return; } finishBlob(context, commandId); }
public static void updateStatus(ComputationContext context, BulkStatus delta) { context.produceRecord(OUTPUT_1, delta.getId(), BulkCodecs.getStatusCodec().encode(delta)); }
@Override public LogRecord<M> read(Duration timeout) throws InterruptedException { if (closed) { throw new IllegalStateException("The tailer has been closed."); } if (records.isEmpty()) { int items = poll(timeout); if (isRebalanced) { isRebalanced = false; log.debug("Rebalance happens during poll, raising exception"); throw new RebalanceException("Partitions has been rebalanced"); } if (items == 0) { if (log.isTraceEnabled()) { log.trace("No data " + id + " after " + timeout.toMillis() + " ms"); } return null; } } ConsumerRecord<String, Bytes> record = records.poll(); lastOffsets.put(new TopicPartition(record.topic(), record.partition()), record.offset()); M value = decodeCodec.decode(record.value().get()); LogPartition partition = LogPartition.of(ns.getLogName(record.topic()), record.partition()); LogOffset offset = new LogOffsetImpl(partition, record.offset()); consumerMoved = false; if (log.isDebugEnabled()) { log.debug(String.format("Read from %s/%s, key: %s, value: %s", offset, group, record.key(), value)); } return new LogRecord<>(value, offset); }
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)); }
@SuppressWarnings("unchecked") protected LogRecord<M> read() { if (closed) { throw new IllegalStateException("The tailer has been closed."); } List<M> value = new ArrayList<>(1); long offset = cqTailer.index(); if (NO_CODEC.equals(codec)) { // default format to keep backward compatibility try { if (!cqTailer.readDocument(w -> value.add((M) w.read(MSG_KEY).object()))) { return null; } } catch (ClassCastException e) { throw new IllegalArgumentException(e); } } else { if (!cqTailer.readDocument(w -> value.add(codec.decode(w.read().bytes())))) { return null; } } return new LogRecord<>(value.get(0), new LogOffsetImpl(partition, offset)); }
/** * 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 void processRecord(ComputationContext context, String inputStream, Record record) { BulkStatus status = codec.decode(record.getData()); if (IndexAction.ACTION_NAME.equals(status.getAction()) && BulkStatus.State.COMPLETED.equals(status.getState())) { logIndexing(status); BulkService bulkService = Framework.getService(BulkService.class); BulkCommand command = bulkService.getCommand(status.getId()); refreshIndexIfNeeded(command); updateAliasIfNeeded(command); } context.askForCheckpoint(); }
@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; }
@Override public void processRecord(ComputationContext context, String inputStream, Record record) { DataBucket in = codec.decode(record.getData()); if (in.getCount() > 0) { BulkRequest bulkRequest = decodeRequest(in); for (DocWriteRequest request : bulkRequest.requests()) { bulkProcessor.add(request); } BulkStatus delta = BulkStatus.deltaOf(in.getCommandId()); delta.setProcessed(in.getCount()); AbstractBulkComputation.updateStatus(context, delta); } updates = true; }
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(); } }
@Override public void processRecord(ComputationContext context, String inputStreamName, Record record) { BulkBucket bucket = BulkCodecs.getBucketCodec().decode(record.getData()); command = getCommand(bucket.getCommandId()); if (command != null) { delta = BulkStatus.deltaOf(command.getId()); delta.setProcessingStartTime(Instant.now()); delta.setProcessed(bucket.getIds().size()); startBucket(record.getKey()); for (List<String> batch : Lists.partition(bucket.getIds(), command.getBatchSize())) { processBatchOfDocuments(batch); } delta.setProcessingEndTime(Instant.now()); endBucket(context, delta); context.askForCheckpoint(); } else { if (isAbortedCommand(bucket.getCommandId())) { log.debug("Skipping aborted command: {}", bucket.getCommandId()); context.askForCheckpoint(); } else { // this requires a manual intervention, the kv store might have been lost log.error("Stopping processing, unknown command: {}, offset: {}, record: {}.", bucket.getCommandId(), context.getLastOffset(), record); context.askForTermination(); } } }
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)); }
protected void processRecord(ComputationContext context, Record record) { BulkCommand command = null; try { command = BulkCodecs.getCommandCodec().decode(record.getData()); String commandId = command.getId(); int bucketSize = command.getBucketSize() > 0 ? command.getBucketSize()