Search in sources :

Example 1 with TaskRelocationPlans

use of com.netflix.titus.grpc.protogen.TaskRelocationPlans in project titus-control-plane by Netflix.

the class TaskRelocationIntegrationTest method findRelocationPlan.

private Optional<TaskRelocationPlan> findRelocationPlan(String taskId) {
    TaskRelocationPlans plans;
    try {
        TestStreamObserver<TaskRelocationPlans> events = new TestStreamObserver<>();
        client.getCurrentTaskRelocationPlans(TaskRelocationQuery.getDefaultInstance(), events);
        plans = events.getLast();
    } catch (Exception e) {
        if (e.getMessage().contains("Relocation workflow not ready yet")) {
            return Optional.empty();
        }
        throw new RuntimeException(e);
    }
    Optional<TaskRelocationPlan> taskPlan = plans.getPlansList().stream().filter(p -> p.getTaskId().equals(taskId)).findFirst();
    if (taskPlan.isPresent()) {
        return taskPlan;
    }
    // Check if already processed
    try {
        TestStreamObserver<TaskRelocationExecution> events = new TestStreamObserver<>();
        client.getTaskRelocationResult(RelocationTaskId.newBuilder().setId(taskId).build(), events);
        return Optional.of(events.getLast().getTaskRelocationPlan());
    } catch (Exception e) {
        Status status = Status.fromThrowable(e);
        if (status.getCode() == Status.Code.NOT_FOUND) {
            return Optional.empty();
        }
        throw new RuntimeException(e);
    }
}
Also used : TestDataFactory(com.netflix.titus.supplementary.relocation.TestDataFactory) TestDataFactory.newSelfManagedDisruptionBudget(com.netflix.titus.supplementary.relocation.TestDataFactory.newSelfManagedDisruptionBudget) Task(com.netflix.titus.api.jobmanager.model.job.Task) Assertions.assertThat(org.assertj.core.api.Assertions.assertThat) TimeoutException(java.util.concurrent.TimeoutException) Supplier(java.util.function.Supplier) TestStreamObserver(com.netflix.titus.testkit.grpc.TestStreamObserver) TitusRuntimes(com.netflix.titus.common.runtime.TitusRuntimes) BatchJobExt(com.netflix.titus.api.jobmanager.model.job.ext.BatchJobExt) RelocationConnectorStubs(com.netflix.titus.supplementary.relocation.RelocationConnectorStubs) Status(io.grpc.Status) ReadOnlyJobOperations(com.netflix.titus.api.jobmanager.service.ReadOnlyJobOperations) TaskRelocationReason(com.netflix.titus.api.relocation.model.TaskRelocationPlan.TaskRelocationReason) TaskRelocationQuery(com.netflix.titus.grpc.protogen.TaskRelocationQuery) Before(org.junit.Before) CollectionsExt.last(com.netflix.titus.common.util.CollectionsExt.last) Job(com.netflix.titus.api.jobmanager.model.job.Job) EventCase(com.netflix.titus.grpc.protogen.RelocationEvent.EventCase) TaskRelocationExecution(com.netflix.titus.grpc.protogen.TaskRelocationExecution) TaskRelocationPlans(com.netflix.titus.grpc.protogen.TaskRelocationPlans) IntegrationTest(com.netflix.titus.testkit.junit.category.IntegrationTest) Test(org.junit.Test) Category(org.junit.experimental.categories.Category) TimeUnit(java.util.concurrent.TimeUnit) TaskRelocationPlan(com.netflix.titus.grpc.protogen.TaskRelocationPlan) Rule(org.junit.Rule) TaskRelocationStatus(com.netflix.titus.grpc.protogen.TaskRelocationStatus) Optional(java.util.Optional) TitusRuntime(com.netflix.titus.common.runtime.TitusRuntime) Clock(com.netflix.titus.common.util.time.Clock) RelocationEvent(com.netflix.titus.grpc.protogen.RelocationEvent) TaskRelocationServiceGrpc(com.netflix.titus.grpc.protogen.TaskRelocationServiceGrpc) RelocationTaskId(com.netflix.titus.grpc.protogen.RelocationTaskId) MockEnvironment(org.springframework.mock.env.MockEnvironment) Status(io.grpc.Status) TaskRelocationStatus(com.netflix.titus.grpc.protogen.TaskRelocationStatus) TaskRelocationExecution(com.netflix.titus.grpc.protogen.TaskRelocationExecution) TestStreamObserver(com.netflix.titus.testkit.grpc.TestStreamObserver) TaskRelocationPlans(com.netflix.titus.grpc.protogen.TaskRelocationPlans) TimeoutException(java.util.concurrent.TimeoutException) TaskRelocationPlan(com.netflix.titus.grpc.protogen.TaskRelocationPlan)

Aggregations

Job (com.netflix.titus.api.jobmanager.model.job.Job)1 Task (com.netflix.titus.api.jobmanager.model.job.Task)1 BatchJobExt (com.netflix.titus.api.jobmanager.model.job.ext.BatchJobExt)1 ReadOnlyJobOperations (com.netflix.titus.api.jobmanager.service.ReadOnlyJobOperations)1 TaskRelocationReason (com.netflix.titus.api.relocation.model.TaskRelocationPlan.TaskRelocationReason)1 TitusRuntime (com.netflix.titus.common.runtime.TitusRuntime)1 TitusRuntimes (com.netflix.titus.common.runtime.TitusRuntimes)1 CollectionsExt.last (com.netflix.titus.common.util.CollectionsExt.last)1 Clock (com.netflix.titus.common.util.time.Clock)1 RelocationEvent (com.netflix.titus.grpc.protogen.RelocationEvent)1 EventCase (com.netflix.titus.grpc.protogen.RelocationEvent.EventCase)1 RelocationTaskId (com.netflix.titus.grpc.protogen.RelocationTaskId)1 TaskRelocationExecution (com.netflix.titus.grpc.protogen.TaskRelocationExecution)1 TaskRelocationPlan (com.netflix.titus.grpc.protogen.TaskRelocationPlan)1 TaskRelocationPlans (com.netflix.titus.grpc.protogen.TaskRelocationPlans)1 TaskRelocationQuery (com.netflix.titus.grpc.protogen.TaskRelocationQuery)1 TaskRelocationServiceGrpc (com.netflix.titus.grpc.protogen.TaskRelocationServiceGrpc)1 TaskRelocationStatus (com.netflix.titus.grpc.protogen.TaskRelocationStatus)1 RelocationConnectorStubs (com.netflix.titus.supplementary.relocation.RelocationConnectorStubs)1 TestDataFactory (com.netflix.titus.supplementary.relocation.TestDataFactory)1