构建Splittable DoFn(SDF)实现GCS对象并行扫描的问题咨询

核心问题：动态前缀的拆分与Checkpoint处理 #

你的核心矛盾是扫描过程中动态生成新前缀，而trySplit会在未收集完所有前缀时触发，导致拆分后的Restriction继续接收新任务，最终checkDone失败。解决思路是把当前正在扫描的前缀和待扫描的前缀列表分离，避免拆分时的状态混乱：

1. Restriction结构拆分
   将ScannerRestriction拆分为两部分：

   • 待扫描的前缀队列（pendingPrefixes）
   • 正在处理的前缀（activePrefix）+ 其分页标记（activePageToken）
     拆分时仅处理待扫描队列，当前正在处理的前缀保留在原Tracker中，避免状态混乱。

2. tryClaim逻辑修正

   • 优先处理当前正在扫描的前缀，直到分页完成（activePageToken为null）
   • 完成当前前缀后，再从待扫描队列中取出下一个前缀开始处理
   • 新发现的子前缀直接追加到待扫描队列，无需通过position传递，避免trySplit时的状态不一致

3. trySplit逻辑优化

   • 仅对待扫描队列进行拆分，当前正在处理的前缀不参与拆分
   • 当待扫描队列长度≤1或存在活跃前缀时，直接返回null，不进行拆分

4. checkDone校验修正
   只有当当前无正在处理的前缀且待扫描队列为空时，才允许通过checkDone，否则抛出异常。

DirectRunner并行执行问题 #

DirectRunner默认单线程执行，开启并行需：

• 在PipelineOptions中设置setDirectRunnerParallelism(N)（N>1）
• 优化trySplit逻辑：当待扫描队列长度>1时，按fractionOfRemainder计算合理的拆分数量，避免返回0。例如队列长度为5、fractionOfRemainder=0.5时，拆分2个前缀到新Restriction。

Splittable DoFn组件生命周期说明 #

单元素并行执行原理 #

Splittable DoFn通过Restriction拆分实现单元素并行：

1. 初始时，每个输入元素对应一个初始Restriction（比如根前缀""）
2. Runner持续调用trySplit将Restriction拆分为多个子Restriction
3. 每个子Restriction分配给独立的Worker线程/进程执行
4. 拆分过程持续到trySplit返回null或达到并行度上限

RestrictionTracker状态修改时机 #

• 初始化：通过@NewTracker创建Tracker，初始状态由@GetInitialRestriction定义
• 运行时：
  • tryClaim：每次调用都会更新Tracker内部状态（如当前处理的前缀、分页标记）
  • trySplit：拆分后，原Tracker保留部分Restriction，新Tracker持有拆分出的剩余部分
• Checkpoint时：Runner调用currentRestriction获取当前状态，序列化后保存到Checkpoint
• 恢复时：从Checkpoint反序列化Restriction，通过@NewTracker创建新Tracker恢复状态

优化后的代码示例 #

@DoFn.BoundedPerElement
private static class ScannerDoFn extends DoFn<String, String> {
    private transient GcsUtil gcsUtil;
    private static final Logger logger = LoggerFactory.getLogger(ScannerDoFn.class);

    @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)) {
                return ProcessContinuation.stop();
            }

            if (position.currentPrefix == null) {
                return ProcessContinuation.resume();
            }

            try {
                Objects objects = gcsUtil.listObjects(
                        bucket,
                        position.currentPrefix,
                        position.currentPageToken,
                        "/");

                if (objects.getItems() != null) {
                    for (StorageObject o : objects.getItems()) {
                        outputReceiver.output(o.getName());
                    }
                }

                if (objects.getPrefixes() != null) {
                    currentRestriction.addPendingPrefixes(objects.getPrefixes());
                }

                position.currentPageToken = objects.getNextPageToken();
                if (position.currentPageToken == null) {
                    position.completedCurrent = true;
                    return ProcessContinuation.resume();
                }
            } catch (Throwable throwable) {
                logger.error("Error scanning prefix {}", position.currentPrefix, throwable);
                position.completedCurrent = true;
                return ProcessContinuation.resume();
            }
        }
    }

    @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 {
    String currentPrefix;
    String currentPageToken;
    boolean completedCurrent;

    public ScannerPosition() {
        this.currentPrefix = null;
        this.currentPageToken = null;
        this.completedCurrent = false;
    }
}

private static class ScannerRestriction {
    final String bucket;
    final LinkedList<String> pendingPrefixes;
    String activePrefix;
    String activePageToken;

    private ScannerRestriction(String bucket) {
        this.bucket = bucket;
        this.pendingPrefixes = Lists.newLinkedList();
        this.activePrefix = null;
        this.activePageToken = null;
    }

    public static ScannerRestriction init(String bucket) {
        ScannerRestriction res = new ScannerRestriction(bucket);
        res.pendingPrefixes.add("");
        return res;
    }

    public ScannerRestriction splitOff(int splitSize) {
        ScannerRestriction split = new ScannerRestriction(bucket);
        for (int i = 0; i < splitSize && !pendingPrefixes.isEmpty(); i++) {
            split.pendingPrefixes.add(pendingPrefixes.poll());
        }
        return split;
    }

    public void addPendingPrefixes(List<String> prefixes) {
        this.pendingPrefixes.addAll(prefixes);
    }

    public boolean hasWork() {
        return activePrefix != null || !pendingPrefixes.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 Coder<List<String>> listCoder = ListCoder.of(StringUtf8Coder.of());
        private final Coder<String> stringCoder = StringUtf8Coder.of();

        @Override
        public void encode(ScannerRestriction value, OutputStream outStream) throws IOException {
            stringCoder.encode(value.bucket, outStream);
            listCoder.encode(value.pendingPrefixes, outStream);
            NullableCoder.of(stringCoder).encode(value.activePrefix, outStream);
            NullableCoder.of(stringCoder).encode(value.activePageToken, outStream);
        }

        @Override
        public ScannerRestriction decode(InputStream inStream) throws IOException {
            String bucket = stringCoder.decode(inStream);
            List<String> pendingPrefixes = listCoder.decode(inStream);
            String activePrefix = NullableCoder.of(stringCoder).decode(inStream);
            String activePageToken = NullableCoder.of(stringCoder).decode(inStream);

            ScannerRestriction res = new ScannerRestriction(bucket);
            res.pendingPrefixes.addAll(pendingPrefixes);
            res.activePrefix = activePrefix;
            res.activePageToken = activePageToken;
            return res;
        }
    }
}

private static class ScannerRestrictionTracker extends RestrictionTracker<ScannerRestriction, ScannerPosition> {
    private final ScannerRestriction restriction;

    ScannerRestrictionTracker(ScannerRestriction restriction) {
        this.restriction = restriction;
    }

    @Override
    public boolean tryClaim(ScannerPosition position) {
        if (position.completedCurrent) {
            restriction.activePrefix = null;
            restriction.activePageToken = null;
            position.completedCurrent = false;
            return tryClaimNextPrefix(position);
        }

        if (restriction.activePrefix != null) {
            position.currentPrefix = restriction.activePrefix;
            position.currentPageToken = restriction.activePageToken;
            return true;
        }

        return tryClaimNextPrefix(position);
    }

    private boolean tryClaimNextPrefix(ScannerPosition position) {
        if (restriction.pendingPrefixes.isEmpty()) {
            position.currentPrefix = null;
            return false;
        }

        restriction.activePrefix = restriction.pendingPrefixes.poll();
        restriction.activePageToken = null;
        position.currentPrefix = restriction.activePrefix;
        position.currentPageToken = null;
        return true;
    }

    @Override
    public ScannerRestriction currentRestriction() {
        return restriction;
    }

    @Override
    public SplitResult<ScannerRestriction> trySplit(double fractionOfRemainder) {
        if (restriction.activePrefix != null) {
            return null;
        }

        int totalPending = restriction.pendingPrefixes.size();
        if (totalPending <= 1) {
            return null;
        }

        int splitSize = (int) Math.round(fractionOfRemainder * totalPending);
        splitSize = Math.max(1, Math.min(splitSize, totalPending - 1));

        ScannerRestriction splitRestriction = restriction.splitOff(splitSize);
        return SplitResult.of(restriction, splitRestriction);
    }

    @Override
    public void checkDone() throws IllegalStateException {
        if (restriction.hasWork()) {
            throw new IllegalStateException("Unfinished work remaining: active prefix or pending prefixes exist");
        }
    }

    @Override
    public IsBounded isBounded() {
        return IsBounded.BOUNDED;
    }
}

内容的提问来源于stack exchange，提问作者Valentin