/**
* Get the final state of each construct used by this task and add it to the {@link
* gobblin.runtime.TaskState}.
*
* @param extractor the {@link gobblin.instrumented.extractor.InstrumentedExtractorBase} used by
* this task.
* @param converter the {@link gobblin.converter.Converter} used by this task.
* @param rowChecker the {@link RowLevelPolicyChecker} used by this task.
*/
private void addConstructsFinalStateToTaskState(
InstrumentedExtractorBase<?, ?> extractor,
Converter<?, ?, ?, ?> converter,
RowLevelPolicyChecker rowChecker) {
ConstructState constructState = new ConstructState();
if (extractor != null) {
constructState.addConstructState(
Constructs.EXTRACTOR, new ConstructState(extractor.getFinalState()));
}
if (converter != null) {
constructState.addConstructState(
Constructs.CONVERTER, new ConstructState(converter.getFinalState()));
}
if (rowChecker != null) {
constructState.addConstructState(
Constructs.ROW_QUALITY_CHECKER, new ConstructState(rowChecker.getFinalState()));
}
int forkIdx = 0;
for (Optional<Fork> fork : this.forks) {
if (fork.isPresent()) {
constructState.addConstructState(
Constructs.FORK_OPERATOR,
new ConstructState(fork.get().getFinalState()),
Integer.toString(forkIdx));
}
forkIdx++;
}
constructState.mergeIntoWorkUnitState(this.taskState);
}
/** Process a (possibly converted) record. */
@SuppressWarnings("unchecked")
private void processRecord(
Object convertedRecord,
ForkOperator forkOperator,
RowLevelPolicyChecker rowChecker,
RowLevelPolicyCheckResults rowResults,
int branches)
throws Exception {
// Skip the record if quality checking fails
if (!rowChecker.executePolicies(convertedRecord, rowResults)) {
return;
}
List<Boolean> forkedRecords = forkOperator.forkDataRecord(this.taskState, convertedRecord);
if (forkedRecords.size() != branches) {
throw new ForkBranchMismatchException(
String.format(
"Number of forked data records [%d] is not equal to number of branches [%d]",
forkedRecords.size(), branches));
}
if (inMultipleBranches(forkedRecords) && !(convertedRecord instanceof Copyable)) {
throw new CopyNotSupportedException(convertedRecord + " is not copyable");
}
// If the record has been successfully put into the queues of every forks
boolean allPutsSucceeded = false;
// Use an array of primitive boolean type to avoid unnecessary boxing/unboxing
boolean[] succeededPuts = new boolean[branches];
// Put the record into the record queue of each fork. A put may timeout and return a false, in
// which
// case the put needs to be retried in the next iteration along with other failed puts. This
// goes on
// until all puts succeed, at which point the task moves to the next record.
while (!allPutsSucceeded) {
allPutsSucceeded = true;
for (int i = 0; i < branches; i++) {
if (succeededPuts[i]) {
continue;
}
if (this.forks.get(i).isPresent() && forkedRecords.get(i)) {
boolean succeeded =
this.forks
.get(i)
.get()
.putRecord(
convertedRecord instanceof Copyable
? ((Copyable) convertedRecord).copy()
: convertedRecord);
succeededPuts[i] = succeeded;
if (!succeeded) {
allPutsSucceeded = false;
}
} else {
succeededPuts[i] = true;
}
}
}
}