Examples with Session org.zalando.nakadi.service.subscription.model.Session used on opensource projects

Example 11 with Session

use of org.zalando.nakadi.service.subscription.model.Session in project nakadi by zalando.

the class SubscriptionControllerTest method whenGetSubscriptionStatThenOk.

@Test
public void whenGetSubscriptionStatThenOk() throws Exception {
    final Subscription subscription = builder().withEventType(TIMELINE.getEventType()).build();
    final Collection<Partition> partitions = Collections.singleton(new Partition(TIMELINE.getEventType(), "0", "xz", null, Partition.State.ASSIGNED));
    final ZkSubscriptionNode zkSubscriptionNode = new ZkSubscriptionNode(partitions, Arrays.asList(new Session("xz", 0)));
    when(subscriptionRepository.getSubscription(subscription.getId())).thenReturn(subscription);
    when(zkSubscriptionClient.getZkSubscriptionNodeLocked()).thenReturn(Optional.of(zkSubscriptionNode));
    final SubscriptionCursorWithoutToken currentOffset = new SubscriptionCursorWithoutToken(TIMELINE.getEventType(), "0", "3");
    final EventTypePartition etp = new EventTypePartition(TIMELINE.getEventType(), "0");
    final Map<EventTypePartition, SubscriptionCursorWithoutToken> offsets = new HashMap<>();
    offsets.put(etp, currentOffset);
    when(zkSubscriptionClient.getOffsets(Collections.singleton(etp))).thenReturn(offsets);
    when(eventTypeRepository.findByName(TIMELINE.getEventType())).thenReturn(EventTypeTestBuilder.builder().name(TIMELINE.getEventType()).build());
    final List<PartitionEndStatistics> statistics = Collections.singletonList(new KafkaPartitionEndStatistics(TIMELINE, 0, 13));
    when(topicRepository.loadTopicEndStatistics(eq(Collections.singletonList(TIMELINE)))).thenReturn(statistics);
    final NakadiCursor currentCursor = mock(NakadiCursor.class);
    when(currentCursor.getEventTypePartition()).thenReturn(new EventTypePartition(TIMELINE.getEventType(), "0"));
    when(cursorConverter.convert((List<SubscriptionCursorWithoutToken>) any())).thenReturn(Collections.singletonList(currentCursor));
    when(cursorOperationsService.calculateDistance(eq(currentCursor), eq(statistics.get(0).getLast()))).thenReturn(10L);
    final List<SubscriptionEventTypeStats> expectedStats = Collections.singletonList(new SubscriptionEventTypeStats(TIMELINE.getEventType(), Collections.singletonList(new SubscriptionEventTypeStats.Partition("0", "assigned", 10L, "xz", AUTO))));
    getSubscriptionStats(subscription.getId()).andExpect(status().isOk()).andExpect(content().string(TestUtils.JSON_TEST_HELPER.matchesObject(new ItemsWrapper<>(expectedStats))));
}

Example 12 with Session

use of org.zalando.nakadi.service.subscription.model.Session in project nakadi by zalando.

the class SubscriptionRebalancerTest method directlyRequestedPartitionsAreCaptured.

@Test
public void directlyRequestedPartitionsAreCaptured() {
    final Partition[] changeset = new SubscriptionRebalancer().apply(ImmutableList.of(new Session("s1", 1), new Session("s2", 1), new Session("s3", 1, ImmutableList.of(new EventTypePartition("et1", "p1"), new EventTypePartition("et1", "p4")))), new Partition[] { new Partition("et1", "p1", "s7", null, ASSIGNED), new Partition("et1", "p2", "s7", null, ASSIGNED), new Partition("et1", "p3", "s7", null, ASSIGNED), new Partition("et1", "p4", "s7", null, ASSIGNED) });
    assertEquals(newHashSet(changeset), newHashSet(new Partition("et1", "p1", "s3", null, ASSIGNED), new Partition("et1", "p2", "s1", null, ASSIGNED), new Partition("et1", "p3", "s2", null, ASSIGNED), new Partition("et1", "p4", "s3", null, ASSIGNED)));
}

Example 13 with Session

use of org.zalando.nakadi.service.subscription.model.Session in project nakadi by zalando.

the class SubscriptionRebalancerTest method directlyAssignedPartitionsAreNotTransferred.

@Test
public void directlyAssignedPartitionsAreNotTransferred() {
    final Partition[] changeset = new SubscriptionRebalancer().apply(ImmutableList.of(new Session("s1", 1, ImmutableList.of(new EventTypePartition("et1", "p1"))), new Session("s2", 1)), new Partition[] { new Partition("et1", "p1", "s1", null, ASSIGNED), new Partition("et1", "p2", null, null, UNASSIGNED) });
    assertEquals(newHashSet(changeset), newHashSet(new Partition("et1", "p2", "s2", null, ASSIGNED)));
}

Example 14 with Session

use of org.zalando.nakadi.service.subscription.model.Session in project nakadi by zalando.

the class SubscriptionRebalancer method rebalanceByWeight.

private Partition[] rebalanceByWeight(final Collection<Session> sessions, final Partition[] currentPartitions) {
    final Map<String, Integer> activeSessionWeights = sessions.stream().collect(Collectors.toMap(Session::getId, Session::getWeight));
    // sorted session ids.
    final List<String> activeSessionIds = activeSessionWeights.keySet().stream().sorted().collect(Collectors.toList());
    // the main part of rebalance - calculate count for each partition.
    final int[] partitionsPerSession = splitByWeight(currentPartitions.length, activeSessionIds.stream().mapToInt(activeSessionWeights::get).toArray());
    // Stage 1. Select partitions that are not assigned to any EXISTING session.
    final List<Partition> toRebalance = Stream.of(currentPartitions).filter(p -> p.mustBeRebalanced(activeSessionIds)).collect(Collectors.toList());
    // State 2. Remove partitions from sessions that have too many of them.
    // 2.1. collect information per session.
    final Map<String, List<Partition>> partitions = Stream.of(currentPartitions).filter(p -> !toRebalance.contains(p)).collect(Collectors.groupingBy(Partition::getEffectiveSession));
    // 2.2. Remove
    for (int idx = 0; idx < activeSessionIds.size(); ++idx) {
        final String sessionId = activeSessionIds.get(idx);
        final int suggestedCount = partitionsPerSession[idx];
        int toTake = (partitions.containsKey(sessionId) ? partitions.get(sessionId).size() : 0) - suggestedCount;
        while (toTake > 0) {
            final List<Partition> candidates = partitions.get(sessionId);
            final Partition toTakeItem = candidates.stream().filter(p -> p.getState() == Partition.State.REASSIGNING).findAny().orElse(candidates.get(candidates.size() - 1));
            candidates.remove(toTakeItem);
            toRebalance.add(toTakeItem);
            toTake -= 1;
        }
    }
    if (!toRebalance.isEmpty()) {
        // 3. Assign partitions to any nodes who are waiting for it.
        final List<Partition> result = new ArrayList<>();
        for (int idx = 0; idx < activeSessionIds.size(); ++idx) {
            final String sessionId = activeSessionIds.get(idx);
            final int suggestedCount = partitionsPerSession[idx];
            final int currentCount = partitions.containsKey(sessionId) ? partitions.get(sessionId).size() : 0;
            for (int i = 0; i < suggestedCount - currentCount; ++i) {
                final Partition toMove = toRebalance.iterator().next();
                toRebalance.remove(toMove);
                result.add(toMove.moveToSessionId(sessionId, activeSessionIds));
            }
        }
        return result.toArray(new Partition[result.size()]);
    } else {
        return new Partition[0];
    }
}