Examples with Batch com.netflix.titus.common.util.rx.batch.Batch used on opensource projects

Example 1 with Batch

use of com.netflix.titus.common.util.rx.batch.Batch in project titus-control-plane by Netflix.

the class DefaultLoadBalancerServiceTest method removeLoadBalancerDeregisterKnownTargets.

@Test
public void removeLoadBalancerDeregisterKnownTargets() {
    String jobId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    JobLoadBalancer jobLoadBalancer = new JobLoadBalancer(jobId, loadBalancerId);
    defaultStubs();
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);
    List<Task> tasks = LoadBalancerTests.buildTasksStarted(5, jobId);
    Collection<LoadBalancerTargetState> expectedTargets = tasks.stream().map(task -> new LoadBalancerTargetState(new LoadBalancerTarget(loadBalancerId, task.getId(), task.getTaskContext().get(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP)), // final state expected after updates are applied
    DEREGISTERED)).collect(Collectors.toList());
    when(v3JobOperations.getTasks(jobId)).thenReturn(CollectionsExt.merge(tasks, LoadBalancerTests.buildTasks(2, jobId, TaskState.StartInitiated), LoadBalancerTests.buildTasks(2, jobId, TaskState.KillInitiated), LoadBalancerTests.buildTasks(3, jobId, TaskState.Finished), LoadBalancerTests.buildTasks(1, jobId, TaskState.Disconnected)));
    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);
    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();
    assertTrue(loadBalancerStore.addOrUpdateLoadBalancer(jobLoadBalancer, JobLoadBalancer.State.ASSOCIATED).await(100, TimeUnit.MILLISECONDS));
    assertTrue(service.removeLoadBalancer(jobId, loadBalancerId).await(100, TimeUnit.MILLISECONDS));
    List<JobLoadBalancerState> jobLoadBalancers = loadBalancerStore.getAssociations().stream().filter(association -> jobId.equals(association.getJobId())).collect(Collectors.toList());
    assertThat(jobLoadBalancers).isNotEmpty();
    assertThat(jobLoadBalancers).hasSize(1);
    JobLoadBalancerState jobLoadBalancerState = jobLoadBalancers.iterator().next();
    assertEquals(loadBalancerId, jobLoadBalancerState.getLoadBalancerId());
    assertEquals(JobLoadBalancer.State.DISSOCIATED, jobLoadBalancerState.getState());
    assertFalse(service.getJobLoadBalancers(jobId).toBlocking().getIterator().hasNext());
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(1);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).containsExactlyInAnyOrderElementsOf(expectedTargets);
    verify(client, never()).registerAll(eq(loadBalancerId), any());
    verify(client).deregisterAll(eq(loadBalancerId), argThat(targets -> targets != null && targets.size() == tasks.size()));
    verifyReconcilerIgnore(loadBalancerId, LoadBalancerTests.ipAddresses(tasks));
}

Example 2 with Batch

use of com.netflix.titus.common.util.rx.batch.Batch in project titus-control-plane by Netflix.

the class DefaultLoadBalancerServiceTest method targetsAreBufferedInBatches.

@Test
public void targetsAreBufferedInBatches() {
    String jobId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    ThreadLocalRandom random = ThreadLocalRandom.current();
    int batchSize = random.nextInt(5, 20);
    defaultStubs();
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);
    List<Task> tasks = LoadBalancerTests.buildTasksStarted(batchSize, jobId);
    when(v3JobOperations.getTasks(jobId)).thenReturn(tasks);
    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);
    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();
    assertTrue(service.addLoadBalancer(jobId, loadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(jobId).toBlocking().first()).isEqualTo(loadBalancerId);
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(1);
    verify(client).registerAll(eq(loadBalancerId), argThat(targets -> targets != null && targets.size() == batchSize));
    verify(client, never()).deregisterAll(eq(loadBalancerId), any());
    verifyReconcilerIgnore(loadBalancerId, LoadBalancerTests.ipAddresses(tasks));
}

Example 3 with Batch

use of com.netflix.titus.common.util.rx.batch.Batch in project titus-control-plane by Netflix.

the class DefaultLoadBalancerServiceTest method newTasksGetRegistered.

@Test
public void newTasksGetRegistered() {
    String jobId = UUID.randomUUID().toString();
    String taskId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();
    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    when(v3JobOperations.getTasks(jobId)).thenReturn(Collections.emptyList());
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);
    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);
    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();
    assertTrue(service.addLoadBalancer(jobId, loadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(jobId).toBlocking().first()).isEqualTo(loadBalancerId);
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    Task launched = ServiceJobTask.newBuilder().withJobId(jobId).withId(taskId).withStatus(TaskStatus.newBuilder().withState(TaskState.Launched).build()).build();
    Task startingWithIp = launched.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.StartInitiated).build()).withTaskContext(CollectionsExt.asMap(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4")).build();
    Task started = startingWithIp.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build()).build();
    LoadBalancerTargetState expectedTarget = new LoadBalancerTargetState(new LoadBalancerTarget(loadBalancerId, started.getId(), "1.2.3.4"), REGISTERED);
    // events with no state transition gets ignored
    taskEvents.onNext(TaskUpdateEvent.newTask(null, launched, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    // events to !Started states get ignored
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, startingWithIp, launched, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    // finally detect the task is UP and gets registered
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, started, startingWithIp, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(1);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).contains(expectedTarget);
    verify(client).registerAll(eq(loadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client, never()).deregisterAll(eq(loadBalancerId), any());
    verifyReconcilerIgnore(loadBalancerId, "1.2.3.4");
}

Example 4 with Batch

use of com.netflix.titus.common.util.rx.batch.Batch in project titus-control-plane by Netflix.

the class DefaultLoadBalancerServiceTest method finishedTasksGetDeregistered.

@Test
public void finishedTasksGetDeregistered() {
    String jobId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();
    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    when(v3JobOperations.getTasks(jobId)).thenReturn(Collections.emptyList());
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);
    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);
    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();
    assertTrue(service.addLoadBalancer(jobId, loadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(jobId).toBlocking().first()).isEqualTo(loadBalancerId);
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    // a task that was prematurely killed before having an IP address associated to it should be ignored
    Task noIp = ServiceJobTask.newBuilder().withJobId(jobId).withId(UUID.randomUUID().toString()).withStatus(TaskStatus.newBuilder().withState(TaskState.KillInitiated).build()).build();
    Task noIpFinished = noIp.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.Finished).build()).build();
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, noIpFinished, noIp, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    // 3 state transitions to 3 different terminal events
    Task first = noIp.toBuilder().withId(UUID.randomUUID().toString()).withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build()).withTaskContext(CollectionsExt.asMap(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.1.1.1")).build();
    Task firstFinished = first.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.Finished).build()).build();
    LoadBalancerTargetState expectedFirstTarget = new LoadBalancerTargetState(new LoadBalancerTarget(loadBalancerId, firstFinished.getId(), "1.1.1.1"), DEREGISTERED);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).doesNotContain(expectedFirstTarget);
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, firstFinished, first, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(1);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).contains(expectedFirstTarget);
    verify(client, never()).registerAll(eq(loadBalancerId), any());
    verify(client).deregisterAll(eq(loadBalancerId), argThat(set -> set.contains("1.1.1.1")));
    verifyReconcilerIgnore(loadBalancerId, "1.1.1.1");
    Task second = first.toBuilder().withId(UUID.randomUUID().toString()).withTaskContext(CollectionsExt.asMap(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "2.2.2.2")).build();
    Task secondKilling = second.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.KillInitiated).build()).build();
    LoadBalancerTargetState expectedSecondTarget = new LoadBalancerTargetState(new LoadBalancerTarget(loadBalancerId, secondKilling.getId(), "2.2.2.2"), DEREGISTERED);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).doesNotContain(expectedSecondTarget);
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, secondKilling, second, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(2);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).contains(expectedSecondTarget);
    verify(client, never()).registerAll(eq(loadBalancerId), any());
    verify(client).deregisterAll(eq(loadBalancerId), argThat(set -> set.contains("2.2.2.2")));
    verifyReconcilerIgnore(loadBalancerId, "2.2.2.2");
    Task third = first.toBuilder().withId(UUID.randomUUID().toString()).withTaskContext(CollectionsExt.asMap(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "3.3.3.3")).build();
    Task thirdDisconnected = third.toBuilder().withStatus(TaskStatus.newBuilder().withState(TaskState.Disconnected).build()).build();
    LoadBalancerTargetState expectedThirdTarget = new LoadBalancerTargetState(new LoadBalancerTarget(loadBalancerId, thirdDisconnected.getId(), "3.3.3.3"), DEREGISTERED);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).doesNotContain(expectedThirdTarget);
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, thirdDisconnected, third, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(3);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block()).contains(expectedThirdTarget);
    verify(client, never()).registerAll(eq(loadBalancerId), any());
    verify(client).deregisterAll(eq(loadBalancerId), argThat(set -> set.contains("3.3.3.3")));
    verifyReconcilerIgnore(loadBalancerId, "3.3.3.3");
}

Example 5 with Batch

use of com.netflix.titus.common.util.rx.batch.Batch in project titus-control-plane by Netflix.

the class DefaultLoadBalancerServiceTest method movedTaskOnlySourceAssociatedWithLoadBalancer.

@Test
public void movedTaskOnlySourceAssociatedWithLoadBalancer() {
    String taskId = UUID.randomUUID().toString();
    String sourceJobId = UUID.randomUUID().toString();
    String targetJobId = UUID.randomUUID().toString();
    String sourceLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();
    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, sourceJobId);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, targetJobId);
    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);
    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();
    assertTrue(service.addLoadBalancer(sourceJobId, sourceLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(sourceJobId).toBlocking().toIterable()).containsExactlyInAnyOrder(sourceLoadBalancerId);
    assertThat(service.getJobLoadBalancers(targetJobId).toBlocking().toIterable()).isEmpty();
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();
    Task moved = ServiceJobTask.newBuilder().withJobId(targetJobId).withId(taskId).withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build()).withTaskContext(CollectionsExt.asMap(TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4", TaskAttributes.TASK_ATTRIBUTES_MOVED_FROM_JOB, sourceJobId)).build();
    // detect the task is moved and gets deregistered from the source
    taskEvents.onNext(TaskUpdateEvent.newTaskFromAnotherJob(null, moved, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    testSubscriber.assertNoErrors().assertValueCount(1);
    verify(client).deregisterAll(eq(sourceLoadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client, never()).registerAll(any(), any());
    verifyReconcilerIgnore(sourceLoadBalancerId, "1.2.3.4");
}