Wait Click Join in YTsaurus Flow (Java)

@Override
public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
    var streamId = message.getStreamId();
    String hitId;
    Long hitTime;
    Hit hit = null;
    Action action = null;

    if ("hit".equals(streamId)) {
        hit = message.getPayload();
        hitId = hit.getHitId();
        hitTime = hit.getHitTime();
    } else if ("action".equals(streamId)) {
        action = message.getPayload();
        hitId = action.getHitId();
        hitTime = action.getHitTime();
    } else {
        throw new IllegalArgumentException("Unknown streamId: %s".formatted(streamId));
    }
    log.debug("Process message (Key: {}, HitId: {}, HitTime: {})", message.getKey(), hitId, hitTime);
    long waitTime = parseDuration(ctx.getComputationParameters().get("wait_for_actions").stringValue()).toSeconds();
    // [BEGIN timer_setup]
    long maxTime = hitTime + waitTime;
    // All messages out of wait range would be dropped.
    if (message.getEventTimestamp() >= maxTime) {
        log.warn("Message is out of wait range, skip (HitId: {}, HitTime: {}, MaxTime: {})",
                hitId, hitTime, maxTime);
        return;
    }
    // Drop late data.
    // [BEGIN late_data_check]
    if (message.getEventTimestamp() < ctx.getEpochInputEventWatermark()) {
        log.warn("Late message, skip (HitId: {}, EventTimestamp: {}, InputEventWatermark: {})",
                hitId, message.getEventTimestamp(), ctx.getEpochInputEventWatermark());
        return;
    }
    // [END late_data_check]
    // [BEGIN state_update]
    ExternalStateAccessor stateAccessor = ctx.getState(JOINED_ACTION_STATE, message);
    PayloadBuilder joinState = stateAccessor.getOrDefault().toBuilder();

    if ("hit".equals(streamId)) {
        log.debug("Process hit message (HitId: {}, HitPayload: {}, Hit: {})", hitId, hit.getHitPayload(), hit);
        joinState.set("hit_payload", hit.getHitPayload());
    } else {
        log.debug("Process action message (HitId: {}, ActionTime: {}, Action: {})",
                hitId, action.getActionTime(), action);
        if (Boolean.TRUE.equals(action.isClick())) {
            joinState.set("click_time", action.getActionTime());
        } else {
            joinState.set("show_time", action.getActionTime());
        }
    }
    log.debug("Set state (KeyHitId: {}, State: {})", hitId, joinState);
    stateAccessor.set(joinState.finish());
    // [END state_update]
    log.debug("Add timer for Key (MaxTime: {}, HitTime: {})", maxTime, hitTime);
    output.addTimer(maxTime, hitTime);
    // [END timer_setup]
}
// [END on_message]

// [BEGIN on_timer]
@Override
public void onTimer(Timer timer, OutputCollector output, RuntimeContext ctx) {
    // Key to pojo conversion is not supported yet.
    String hitId = timer.getKey().get("hit_id", String.class);
    log.debug("Process timer (KeyHitId: {}, Timer: {})", hitId, timer);
    ExternalStateAccessor stateAccessor = ctx.getState(JOINED_ACTION_STATE, timer);
    Payload joinState = stateAccessor.get();
    if (joinState == null) {
        throw new IllegalStateException("Empty state for timer: " + timer);
    }
    log.debug("Process state (KeyHitId: {}, State: {})", hitId, joinState);
    if ("timer".equals(timer.getStreamId())) {
        // ExternalState to pojo conversion is not supported yet.
        if (joinState.get("show_time", Long.class) != null
                && joinState.get("hit_payload", String.class) != null
                && joinState.get("show_time", Long.class) != 0) {
            log.debug("Click join: {}", joinState);
            JoinedAction joinedAction = new JoinedAction();
            long hitTime = timer.getKey().get("hit_time", Long.class);
            joinedAction.setHitId(hitId);
            joinedAction.setHitTime(hitTime);
            joinedAction.setShowTime(joinState.get("show_time", Long.class));
            if (joinState.get("click_time", Long.class) != null && joinState.get("click_time", Long.class) != 0) {
                joinedAction.setClick(true);
                joinedAction.setClickTime(joinState.get("click_time", Long.class));
            } else {
                joinedAction.setClick(false);
                joinedAction.setClickTime(0L);
            }
            joinedAction.setHitPayload(joinState.get("hit_payload", String.class));
            output.addMessage(Message.builder().setStreamId("joined_action").setPayload(joinedAction).build());
        }
        stateAccessor.clear();
    } else {
        throw new IllegalArgumentException("Unknown streamId: %s".formatted(timer.getStreamId()));
    }
}

if (message.getEventTimestamp() < ctx.getEpochInputEventWatermark()) {
    log.warn("Late message, skip (HitId: {}, EventTimestamp: {}, InputEventWatermark: {})",
            hitId, message.getEventTimestamp(), ctx.getEpochInputEventWatermark());
    return;
}

long maxTime = hitTime + waitTime;
// All messages out of wait range would be dropped.
if (message.getEventTimestamp() >= maxTime) {
    log.warn("Message is out of wait range, skip (HitId: {}, HitTime: {}, MaxTime: {})",
            hitId, hitTime, maxTime);
    return;
}
// Drop late data.
// [BEGIN late_data_check]
if (message.getEventTimestamp() < ctx.getEpochInputEventWatermark()) {
    log.warn("Late message, skip (HitId: {}, EventTimestamp: {}, InputEventWatermark: {})",
            hitId, message.getEventTimestamp(), ctx.getEpochInputEventWatermark());
    return;
}
// [END late_data_check]
// [BEGIN state_update]
ExternalStateAccessor stateAccessor = ctx.getState(JOINED_ACTION_STATE, message);
PayloadBuilder joinState = stateAccessor.getOrDefault().toBuilder();

if ("hit".equals(streamId)) {
    log.debug("Process hit message (HitId: {}, HitPayload: {}, Hit: {})", hitId, hit.getHitPayload(), hit);
    joinState.set("hit_payload", hit.getHitPayload());
} else {
    log.debug("Process action message (HitId: {}, ActionTime: {}, Action: {})",
            hitId, action.getActionTime(), action);
    if (Boolean.TRUE.equals(action.isClick())) {
        joinState.set("click_time", action.getActionTime());
    } else {
        joinState.set("show_time", action.getActionTime());
    }
}
log.debug("Set state (KeyHitId: {}, State: {})", hitId, joinState);
stateAccessor.set(joinState.finish());
// [END state_update]
log.debug("Add timer for Key (MaxTime: {}, HitTime: {})", maxTime, hitTime);
output.addTimer(maxTime, hitTime);

ExternalStateAccessor stateAccessor = ctx.getState(JOINED_ACTION_STATE, message);
PayloadBuilder joinState = stateAccessor.getOrDefault().toBuilder();

if ("hit".equals(streamId)) {
    log.debug("Process hit message (HitId: {}, HitPayload: {}, Hit: {})", hitId, hit.getHitPayload(), hit);
    joinState.set("hit_payload", hit.getHitPayload());
} else {
    log.debug("Process action message (HitId: {}, ActionTime: {}, Action: {})",
            hitId, action.getActionTime(), action);
    if (Boolean.TRUE.equals(action.isClick())) {
        joinState.set("click_time", action.getActionTime());
    } else {
        joinState.set("show_time", action.getActionTime());
    }
}
log.debug("Set state (KeyHitId: {}, State: {})", hitId, joinState);
stateAccessor.set(joinState.finish());

// The stream schema is derived from this class: the runner writes it into the pipeline spec,
// so "hit" needs no schema of its own there. Field order defines column order, and
// columnDefinition pins a column type where the default for the field type differs.
@Entity
public class Hit {
    @Column(name = "hit_id", columnDefinition = "string")
    private String hitId;
    @Column(name = "hit_time", columnDefinition = "uint64")
    private Long hitTime;
    @Column(name = "hit_payload", columnDefinition = "string")
    private String hitPayload;

    public Hit() {
    }

    public Hit(String hitId, Long hitTime, String hitPayload) {
        this.hitId = hitId;
        this.hitTime = hitTime;
        this.hitPayload = hitPayload;
    }

    public String getHitId() {
        return hitId;
    }

    public void setHitId(String hitId) {
        this.hitId = hitId;
    }

    public Long getHitTime() {
        return hitTime;
    }

    public void setHitTime(Long hitTime) {
        this.hitTime = hitTime;
    }

    public String getHitPayload() {
        return hitPayload;
    }

    public void setHitPayload(String hitPayload) {
        this.hitPayload = hitPayload;
    }

    @Override
    public String toString() {
        return "Hit{" +
                "hitId='" + hitId + '\'' +
                ", hitTime=" + hitTime +
                ", hitPayload='" + hitPayload + '\'' +
                '}';
    }
}

// The stream schema is derived from this class: the runner writes it into the pipeline spec,
// so "action" needs no schema of its own there. Field order defines column order, and
// columnDefinition pins a column type where the default for the field type differs.
@Entity
public class Action {
    @Column(name = "hit_id", columnDefinition = "string")
    private String hitId;
    @Column(name = "hit_time", columnDefinition = "uint64")
    private Long hitTime;
    @Column(name = "is_click")
    private Boolean click;
    @Column(name = "action_time", columnDefinition = "uint64")
    private Long actionTime;

    public Action() {
    }

    public Action(String hitId, Long hitTime, Long actionTime, Boolean isClick) {
        this.hitId = hitId;
        this.hitTime = hitTime;
        this.actionTime = actionTime;
        this.click = isClick;
    }

    public String getHitId() {
        return hitId;
    }

    public void setHitId(String hitId) {
        this.hitId = hitId;
    }

    public Long getHitTime() {
        return hitTime;
    }

    public void setHitTime(Long hitTime) {
        this.hitTime = hitTime;
    }

    public Long getActionTime() {
        return actionTime;
    }

    public void setActionTime(Long actionTime) {
        this.actionTime = actionTime;
    }

    public Boolean isClick() {
        return click;
    }

    public void setIsClick(Boolean click) {
        this.click = click;
    }

    @Override
    public String toString() {
        return "Action{" +
                "hitId='" + hitId + '\'' +
                ", hitTime=" + hitTime +
                ", actionTime=" + actionTime +
                ", click=" + click +
                '}';
    }
}

// The stream schema is derived from this class: the runner writes it into the pipeline spec,
// so "joined_action" needs no schema of its own there. Field order defines column order, and
// columnDefinition pins a column type where the default for the field type differs.
@Entity
public class JoinedAction {
    @Column(name = "hit_id", columnDefinition = "string")
    private String hitId;
    @Column(name = "hit_time", columnDefinition = "uint64")
    private Long hitTime;
    @Column(name = "is_click")
    private Boolean isClick;
    @Column(name = "show_time", columnDefinition = "uint64")
    private Long showTime;
    @Column(name = "click_time", columnDefinition = "uint64")
    private Long clickTime;
    @Column(name = "hit_payload", columnDefinition = "string")
    private String hitPayload;

    public JoinedAction() {
    }

    public JoinedAction(String hitId, Long hitTime, Boolean isClick, Long showTime, Long clickTime, String hitPayload) {
        this.hitId = hitId;
        this.hitTime = hitTime;
        this.isClick = isClick;
        this.showTime = showTime;
        this.clickTime = clickTime;
        this.hitPayload = hitPayload;
    }

    public String getHitId() {
        return hitId;
    }

    public void setHitId(String hitId) {
        this.hitId = hitId;
    }

    public Long getHitTime() {
        return hitTime;
    }

    public void setHitTime(Long hitTime) {
        this.hitTime = hitTime;
    }

    public Boolean getClick() {
        return isClick;
    }

    public void setClick(Boolean click) {
        isClick = click;
    }

    public Long getShowTime() {
        return showTime;
    }

    public void setShowTime(Long showTime) {
        this.showTime = showTime;
    }

    public Long getClickTime() {
        return clickTime;
    }

    public void setClickTime(Long clickTime) {
        this.clickTime = clickTime;
    }

    public String getHitPayload() {
        return hitPayload;
    }

    public void setHitPayload(String hitPayload) {
        this.hitPayload = hitPayload;
    }
}

@FlowComputation(id = "join")

@Configuration
public class PipelineConfiguration implements ComputationProvider {

    @Override
    public List<FlowStream<?>> getStreams() {
        return List.of(
                FlowStreams.typed("hit", Hit.class),
                FlowStreams.typed("action", Action.class),
                FlowStreams.typed("joined_action", JoinedAction.class)
        );
    }
}

public static void main(String[] args) throws Exception {
    new SpringApplicationBuilder(PipelineMain.class)
            .run(args);
}

    @BeforeEach
    void init() throws IOException {
        var pipelineContext = new PipelineContext();

        Computation join = Computation.builder()
                .setComputationId("join")
                .setProcessFunction(new JoinProcessFunction())
                .build();
        pipelineContext.registerComputation(join);

        pipelineContext.registerStream(FlowStreams.typed("hit", Hit.class));
        pipelineContext.registerStream(FlowStreams.typed("action", Action.class));
        pipelineContext.registerStream(FlowStreams.typed("joined_action", JoinedAction.class));

        // External state schema for "/join-state"
        this.joinStateSchema = TableSchema.builder()
                .addValue("hit_payload", ColumnValueType.STRING)
                .addValue("show_time", ColumnValueType.UINT64)
                .addValue("click_time", ColumnValueType.UINT64)
                .build();
        // The pipeline spec is packaged as a classpath resource in the app module's
        // src/main/resources, so the test reads the single source of truth directly.
        try (var specStream = getClass().getResourceAsStream("/pipeline.yson")) {
            YTreeNode spec = YTreeTextSerializer.deserialize(specStream);
            this.testHarness = TestComputationHarness.builder()
                    .setPipelineContext(pipelineContext)
                    .setPipelineSpec(spec)
                    .addExternalStateSchema("/join-state", joinStateSchema)
                    .build();
        }
    }

    @Test
    void testHitMessageStoresHitPayloadInState() {
        String hitId = "hit-1";
        long hitTime = BASE_HIT_TIME;

        var messages = List.of(buildHitMessage(hitId, hitTime, "payload-data"));

        var request = TestDoProcessRequest.builder("join")
                .setMessages(messages)
                .setWatermarks(defaultWatermarks())
                .build();

        var response = testHarness.doProcess(request);

        // No output messages yet — timer is scheduled, output comes on timer fire
        assertTrue(response.getOutputMessagesFlatten().isEmpty());

        // A timer must be scheduled
        assertEquals(1, response.getOutputTimersFlatten().size());
        var timer = response.getOutputTimersFlatten().getFirst();
        assertEquals(hitTime + WAIT_SECONDS, timer.getTriggerTimestamp());

        // External state must be updated with hit_payload
        assertEquals(1, response.allStates().externalSize("/join-state"));
        Payload key = buildKey(hitId, hitTime);
        var state = response.allStates().get(JoinProcessFunction.JOINED_ACTION_STATE, key).get();
        assertEquals("payload-data", state.get("hit_payload", String.class));
    }

    @Test
    void testFullJoinFlow_HitThenShowThenClickThenTimer() {
        String hitId = "hit-20";
        long hitTime = BASE_HIT_TIME;
        long showTime = hitTime + 3L;
        long clickTime = hitTime + 7L;

        // Step 1: process hit message
        var hitMessages = List.of(buildHitMessage(hitId, hitTime, "full-payload"));
        var hitRequest = TestDoProcessRequest.builder("join")
                .setMessages(hitMessages)
                .setWatermarks(defaultWatermarks())
                .build();
        var hitResponse = testHarness.doProcess(hitRequest);

        assertTrue(hitResponse.getOutputMessagesFlatten().isEmpty());
        assertEquals(1, hitResponse.getOutputTimersFlatten().size());
        assertEquals(1, hitResponse.allStates().externalSize("/join-state"));

        Payload key = buildKey(hitId, hitTime);
        Payload stateAfterHit = hitResponse.allStates()
                .get(JoinProcessFunction.JOINED_ACTION_STATE, key).get();
        assertEquals("full-payload", stateAfterHit.get("hit_payload", String.class));

        // Step 2: process show action message, passing the state from step 1
        var showMessages = List.of(buildActionMessage(hitId, hitTime, showTime, false, hitTime + 1L));
        var showRequest = TestDoProcessRequest.builder("join")
                .setMessages(showMessages)
                .setState(JoinProcessFunction.JOINED_ACTION_STATE, key, stateAfterHit)
                .setWatermarks(defaultWatermarks())
                .build();
        var showResponse = testHarness.doProcess(showRequest);

        assertTrue(showResponse.getOutputMessagesFlatten().isEmpty());
        Payload stateAfterShow = showResponse.allStates().get(JoinProcessFunction.JOINED_ACTION_STATE, key)
                .get();
        assertEquals("full-payload", stateAfterShow.get("hit_payload", String.class));
        assertEquals(showTime, stateAfterShow.get("show_time", Long.class));

        // Step 3: process click action message, passing the state from step 2
        var clickMessages = List.of(buildActionMessage(hitId, hitTime, clickTime, true, hitTime + 2L));
        var clickRequest = TestDoProcessRequest.builder("join")
                .setMessages(clickMessages)
                .setState(JoinProcessFunction.JOINED_ACTION_STATE, key, stateAfterShow)
                .setWatermarks(defaultWatermarks())
                .build();
        var clickResponse = testHarness.doProcess(clickRequest);

        assertTrue(clickResponse.getOutputMessagesFlatten().isEmpty());
        Payload stateAfterClick = clickResponse.allStates().get(JoinProcessFunction.JOINED_ACTION_STATE, key)
                .get();
        assertEquals(clickTime, stateAfterClick.get("click_time", Long.class));

        // Step 4: fire the timer with the accumulated state
        Timer timer = buildTimer(hitId, hitTime);
        var timerRequest = TestDoProcessRequest.builder("join")
                .setTimers(List.of(timer))
                .setState(JoinProcessFunction.JOINED_ACTION_STATE, key, stateAfterClick)
                .build();
        var timerResponse = testHarness.doProcess(timerRequest);

        assertEquals(1, timerResponse.getOutputMessagesFlatten().size());
        JoinedAction joinedAction = timerResponse.getOutputMessagesFlatten().getFirst().getPayload();
        assertEquals(hitId, joinedAction.getHitId());
        assertEquals(hitTime, joinedAction.getHitTime());
        assertEquals("full-payload", joinedAction.getHitPayload());
        assertEquals(showTime, joinedAction.getShowTime());
        assertTrue(joinedAction.getClick());
        assertEquals(clickTime, joinedAction.getClickTime());

        assertNull(timerResponse.allStates().get(JoinProcessFunction.JOINED_ACTION_STATE, key).get());
    }