Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

assign all splits correctly to source readers #37

Merged
merged 3 commits into from
Jan 27, 2025
Merged

assign all splits correctly to source readers #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d0d9975
assign splits correctly to source readers
tilakraj94 Jan 16, 2025
3ea2253
add null check for next splits
tilakraj94 Jan 16, 2025
b059460
assign splits properly to source readers
tilakraj94 Jan 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 70 additions & 9 deletions src/main/java/io/synadia/flink/v0/enumerator/NatsSourceEnumerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,25 @@
import io.synadia.flink.v0.source.split.NatsSubjectSplit;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.connector.source.SplitsAssignment;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import java.util.ArrayDeque;
import java.util.Collection;
import java.util.List;
import java.util.Queue;
import java.util.*;

import static io.synadia.flink.v0.utils.MiscUtils.generatePrefixedId;
import static org.apache.flink.util.Preconditions.checkNotNull;

public class NatsSourceEnumerator implements SplitEnumerator<NatsSubjectSplit, Collection<NatsSubjectSplit>> {
private static final Logger LOG = LoggerFactory.getLogger(NatsSourceEnumerator.class);

private final String id;
private final SplitEnumeratorContext<NatsSubjectSplit> context;
private final Queue<NatsSubjectSplit> remainingSplits;
private List<List<NatsSubjectSplit>> precomputedSplitAssignments;



public NatsSourceEnumerator(String sourceId,
SplitEnumeratorContext<NatsSubjectSplit> context,
Expand All @@ -28,24 +33,80 @@ public NatsSourceEnumerator(String sourceId,
id = generatePrefixedId(sourceId);
this.context = checkNotNull(context);
this.remainingSplits = splits == null ? new ArrayDeque<>() : new ArrayDeque<>(splits);
this.precomputedSplitAssignments = Collections.synchronizedList(new LinkedList<>());
}

@Override
public void start() {
int totalSplits = remainingSplits.size();
int parallelism = context.currentParallelism();

// Calculate the minimum splits per reader and leftover splits
int minimumSplitsPerReader = totalSplits / parallelism;
int leftoverSplits = totalSplits % parallelism;

// Precompute split assignments
List<List<NatsSubjectSplit>>splitAssignments = preComputeSplitsAssignments(parallelism, minimumSplitsPerReader, leftoverSplits);

// Store precomputed split assignments
this.precomputedSplitAssignments = splitAssignments;
LOG.debug("{} | Precomputed split assignments: {}", id, splitAssignments);
}

private List<List<NatsSubjectSplit>> preComputeSplitsAssignments (int parallelism, int minimumSplitsPerReader, int leftoverSplits) {
List<List<NatsSubjectSplit>> splitAssignments = new ArrayList<>(parallelism);

// Initialize lists
for (int i = 0; i < parallelism; i++) {
splitAssignments.add(new ArrayList<>());
}

// Distribute splits evenly among subtasks
for (int j = 0; j < parallelism; j++) {
List<NatsSubjectSplit> readerSplits = splitAssignments.get(j);

// Assign minimum splits to each reader
for (int i = 0; i < minimumSplitsPerReader && !remainingSplits.isEmpty(); i++) {
readerSplits.add(remainingSplits.poll());
}

// Assign one leftover split if available
if (leftoverSplits > 0 && !remainingSplits.isEmpty()) {
readerSplits.add(remainingSplits.poll());
leftoverSplits--;
}
}

return splitAssignments;
}

@Override
public void close() {
// remove precomputed split assignments if any
precomputedSplitAssignments.clear();
}

@Override
public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
final NatsSubjectSplit nextSplit = remainingSplits.poll();
if (nextSplit != null) {
context.assignSplit(nextSplit, subtaskId);
}
else {
int size = precomputedSplitAssignments.size();

if (size == 0) {
LOG.debug("{} | No more splits available for subtask {}", id, subtaskId);
context.signalNoMoreSplits(subtaskId);
} else {
// O(1) operation with LinkedList
// Remove the first element from the list
// and assign splits to subtask
List<NatsSubjectSplit> splits = precomputedSplitAssignments.remove(0);
if (splits.isEmpty()) {
LOG.debug("{} | Empty split assignment for subtask {}", id, subtaskId);
context.signalNoMoreSplits(subtaskId);
} else {

// Assign splits to subtask
LOG.debug("{} | Assigning splits {} to subtask {}", id, splits, subtaskId);
context.assignSplits(new SplitsAssignment<>(Map.of(subtaskId, splits)));
}
}
}

Expand Down
108 changes: 108 additions & 0 deletions src/test/java/io/synadia/io/synadia/flink/v0/NatsSourceEnumeratorTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
package io.synadia.io.synadia.flink.v0;

import io.synadia.flink.v0.enumerator.NatsSourceEnumerator;
import io.synadia.flink.v0.source.split.NatsSubjectSplit;
import io.synadia.io.synadia.flink.TestBase;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.connector.source.SplitsAssignment;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import org.mockito.ArgumentCaptor;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

class NatsSourceEnumeratorTests extends TestBase {

private SplitEnumeratorContext<NatsSubjectSplit> context;
private Queue<NatsSubjectSplit> splitsQueue;
private NatsSourceEnumerator enumerator;

@BeforeEach
@SuppressWarnings("unchecked")
void setup() {
context = mock(SplitEnumeratorContext.class);
splitsQueue = new ArrayDeque<>();
}

static Stream<TestParameters> provideTestParameters() {
return Stream.of(
new TestParameters(5, 3, generateSplits(3), "Splits < Parallelism"),
new TestParameters(3, 5, generateSplits(5), "Splits > Parallelism"),
new TestParameters(3, 3, generateSplits(3), "Splits == Parallelism"),
new TestParameters(89, 9, generateSplits(9), "High Parallelism with few splits, Splits <<<< Parallelism"),
new TestParameters(5, 100, generateSplits(100), "More Splits with less parallelism, Splits >>>>>> Parallelism")
);
}

@DisplayName("NatsSourceEnumerator Parameterized Test")
@ParameterizedTest(name = "{index} - {0}")
@MethodSource("provideTestParameters")
void testHandleSplitRequest(TestParameters params) {
// Setup splits queue
splitsQueue.addAll(params.splits.stream().map(NatsSubjectSplit::new).collect(Collectors.toList()));

when(context.currentParallelism()).thenReturn(params.parallelism);
enumerator = new NatsSourceEnumerator("test", context, splitsQueue);

// precompute split assignments
enumerator.start();

for (int subtaskId = 0; subtaskId < params.parallelism; subtaskId++) {
enumerator.handleSplitRequest(subtaskId, null);
}

// Capture arguments passed to assignSplits
ArgumentCaptor<SplitsAssignment<NatsSubjectSplit>> captor = ArgumentCaptor.forClass(SplitsAssignment.class);
verify(context, times(params.expectedInvocations)).assignSplits(captor.capture());

List<SplitsAssignment<NatsSubjectSplit>> assignments = captor.getAllValues();
assertEquals(params.expectedInvocations, assignments.size());

// Verify splits assigned
List<String> assignedSplits = new ArrayList<>();
for (SplitsAssignment<NatsSubjectSplit> assignment : assignments) {
assignment.assignment().values().forEach(splits -> {
assertTrue(splits.size() <= params.splits.size() / params.parallelism + 1, "Each subtask should get the correct number of splits");
splits.forEach(split -> assignedSplits.add(split.getSubject()));
});
}

assertEquals(params.splits, assignedSplits);
}

static List<String> generateSplits(int count) {
List<String> splits = new ArrayList<>();
for (int i = 0; i < count; i++) {
splits.add("split" + i);
}
return splits;
}

static class TestParameters {
final int parallelism;
final int splitsCount;
final List<String> splits;
final int expectedInvocations;
final String description;

TestParameters(int parallelism, int splitsCount, List<String> splits, String description) {
this.parallelism = parallelism;
this.splitsCount = splitsCount;
this.splits = splits;
this.expectedInvocations = Math.min(parallelism, splitsCount);
this.description = description;
}

@Override
public String toString() {
return description;
}
}
}
Loading