Building Splittable DoFn (SDF) to list objects in GCS - google-cloud-dataflow

Question

I am trying to develop custom source for parallel GCS content scanning. The naive approach would be to loop through listObjects function calls:
while (...) {
Objects objects = gcsUtil.listObjects(bucket, prefix, pageToken);
pageToken = objects.getNextPageToken();
...
}
The problem is performance for the tens of millions objects.
Instead of the single thread code we can add delimiter / and submit parallel processed for each prefix found:
...
Objects objects = gcsUtil.listObjects(bucket, prefix, pageToken, "/");
for (String subPrefix : object.getPrefixes()) {
scanAsync(bucket, subPrefix);
}
...
Next idea was to try to wrap this logic in Splittable DoFn.
Choice of RestrictionTracker: I don't see how can be used any of exiting RestrictionTracker. So decided to write own. Restriction itself is basically list of prefixes to scan. tryClaim checks if there is more prefixes left and receive newly scanned to append them to current restriction. trySplit splits list of prefixes.
The problem that I faced that trySplit can be called before all subPrefixes are found. In this case current restriction may receive more work to do after it was splitted. And it seems that trySplit is being called until it returns a not null value for a given RestrictionTracker: after certain number of elements goes to the output or after 1 second via scheduler or when ProcessContext.resume() returned. This doesn't work in my case as new prefixes can be found at any time. And I can't checkpoint via return ProcessContext.resume() because if split was already done, possible work that left in current restriction will cause checkDone() function to fail.
Another problem that I suspect that I couldn't achieve parallel execution in DirectRunner. As trySplit was always called with fractionOfRemainder=0 and new RestrictionTracker was started only after current one completed its piece of work.
It would be also great to read more detailed explanation about Splittable DoFn components lifecycle. How parallel execution per element is achieved. And how and when state of RestrictionTracker can be modified.
UPD: Adding simplified code that should show intended implementation
#DoFn.BoundedPerElement
private static class ScannerDoFn extends DoFn<String, String> {
private transient GcsUtil gcsUtil;
#GetInitialRestriction
public ScannerRestriction getInitialRestriction(#Element String bucket) {
return ScannerRestriction.init(bucket);
}
#ProcessElement
public ProcessContinuation processElement(
ProcessContext c,
#Element String bucket,
RestrictionTracker<ScannerRestriction, ScannerPosition> tracker,
OutputReceiver<String> outputReceiver) {
if (gcsUtil == null) {
gcsUtil = c.getPipelineOptions().as(GcsOptions.class).getGcsUtil();
}
ScannerRestriction currentRestriction = tracker.currentRestriction();
ScannerPosition position = new ScannerPosition();
while (true) {
if (tracker.tryClaim(position)) {
if (position.completedCurrent) {
// position.clear();
// ideally I would to get checkpoint here before starting new work
return ProcessContinuation.resume();
}
try {
Objects objects = gcsUtil.listObjects(
currentRestriction.bucket,
position.currentPrefix,
position.currentPageToken,
"/");
if (objects.getItems() != null) {
for (StorageObject o : objects.getItems()) {
outputReceiver.output(o.getName());
}
}
if (objects.getPrefixes() != null) {
position.newPrefixes.addAll(objects.getPrefixes());
}
position.currentPageToken = objects.getNextPageToken();
if (position.currentPageToken == null) {
position.completedCurrent = true;
}
} catch (Throwable throwable) {
logger.error("Error during scan", throwable);
}
} else {
return ProcessContinuation.stop();
}
}
}
#NewTracker
public RestrictionTracker<ScannerRestriction, ScannerPosition> restrictionTracker(#Restriction ScannerRestriction restriction) {
return new ScannerRestrictionTracker(restriction);
}
#GetRestrictionCoder
public Coder<ScannerRestriction> getRestrictionCoder() {
return ScannerRestriction.getCoder();
}
}
public static class ScannerPosition {
private String currentPrefix;
private String currentPageToken;
private final List<String> newPrefixes;
private boolean completedCurrent;
public ScannerPosition() {
this.currentPrefix = null;
this.currentPageToken = null;
this.newPrefixes = Lists.newArrayList();
this.completedCurrent = false;
}
public void clear() {
this.currentPageToken = null;
this.currentPrefix = null;
this.completedCurrent = false;
}
}
private static class ScannerRestriction {
private final String bucket;
private final LinkedList<String> prefixes;
private ScannerRestriction(String bucket) {
this.bucket = bucket;
this.prefixes = Lists.newLinkedList();
}
public static ScannerRestriction init(String bucket) {
ScannerRestriction res = new ScannerRestriction(bucket);
res.prefixes.add("");
return res;
}
public ScannerRestriction empty() {
return new ScannerRestriction(bucket);
}
public boolean isEmpty() {
return prefixes.isEmpty();
}
public static Coder<ScannerRestriction> getCoder() {
return ScannerRestrictionCoder.INSTANCE;
}
private static class ScannerRestrictionCoder extends AtomicCoder<ScannerRestriction> {
private static final ScannerRestrictionCoder INSTANCE = new ScannerRestrictionCoder();
private final static Coder<List<String>> listCoder = ListCoder.of(StringUtf8Coder.of());
#Override
public void encode(ScannerRestriction value, OutputStream outStream) throws IOException {
NullableCoder.of(StringUtf8Coder.of()).encode(value.bucket, outStream);
listCoder.encode(value.prefixes, outStream);
}
#Override
public ScannerRestriction decode(InputStream inStream) throws IOException {
String bucket = NullableCoder.of(StringUtf8Coder.of()).decode(inStream);
List<String> prefixes = listCoder.decode(inStream);
ScannerRestriction res = new ScannerRestriction(bucket);
res.prefixes.addAll(prefixes);
return res;
}
}
}
private static class ScannerRestrictionTracker extends RestrictionTracker<ScannerRestriction, ScannerPosition> {
private ScannerRestriction restriction;
private ScannerPosition lastPosition = null;
ScannerRestrictionTracker(ScannerRestriction restriction) {
this.restriction = restriction;
}
#Override
public boolean tryClaim(ScannerPosition position) {
restriction.prefixes.addAll(position.newPrefixes);
position.newPrefixes.clear();
if (position.completedCurrent) {
// completed work for current prefix
assert lastPosition != null && lastPosition.currentPrefix.equals(position.currentPrefix);
lastPosition = null;
return true; // return true but we would need to claim again if we need to get next prefix
} else if (lastPosition != null && lastPosition.currentPrefix.equals(position.currentPrefix)) {
// proceed work for current prefix
lastPosition = position;
return true;
}
// looking for next prefix
assert position.currentPrefix == null;
assert lastPosition == null;
if (restriction.isEmpty()) {
// no work to do
return false;
}
position.currentPrefix = restriction.prefixes.poll();
lastPosition = position;
return true;
}
#Override
public ScannerRestriction currentRestriction() {
return restriction;
}
#Override
public SplitResult<ScannerRestriction> trySplit(double fractionOfRemainder) {
if (lastPosition == null && restriction.isEmpty()) {
// no work at all
return null;
}
if (lastPosition != null && restriction.isEmpty()) {
// work at the moment only at currently scanned prefix
return SplitResult.of(restriction, restriction.empty());
}
int size = restriction.prefixes.size();
int newSize = new Double(Math.round(fractionOfRemainder * size)).intValue();
if (newSize == 0) {
ScannerRestriction residual = restriction;
restriction = restriction.empty();
return SplitResult.of(restriction, residual);
}
ScannerRestriction residual = restriction.empty();
for (int i=newSize; i<=size; i++) {
residual.prefixes.add(restriction.prefixes.removeLast());
}
return SplitResult.of(restriction, residual);
}
#Override
public void checkDone() throws IllegalStateException {
if (lastPosition != null) {
throw new IllegalStateException("Called checkDone on not completed job");
}
}
#Override
public IsBounded isBounded() {
return IsBounded.UNBOUNDED;
}
}