com.google.longrunning.Operation Java Examples

public RedisShardBackplane(
    RedisShardBackplaneConfig config,
    String source,
    Function<Operation, Operation> onPublish,
    Function<Operation, Operation> onComplete,
    Predicate<Operation> isPrequeued,
    Predicate<Operation> isDispatched)
    throws ConfigurationException {
  this(
      config,
      source,
      onPublish,
      onComplete,
      isPrequeued,
      isDispatched,
      JedisClusterFactory.create(config));
}
 

static void listDatabaseOperations(
    InstanceAdminClient instanceAdminClient,
    DatabaseAdminClient dbAdminClient,
    InstanceId instanceId) {
  Instance instance = instanceAdminClient.getInstance(instanceId.getInstance());
  // Get optimize restored database operations.
  String filter = "(metadata.@type:type.googleapis.com/"
                  + "google.spanner.admin.database.v1.OptimizeRestoredDatabaseMetadata)";
  for (Operation op : instance.listDatabaseOperations(Options.filter(filter)).iterateAll()) {
    try {
      OptimizeRestoredDatabaseMetadata metadata =
          op.getMetadata().unpack(OptimizeRestoredDatabaseMetadata.class);
      System.out.println(String.format(
          "Database %s restored from backup is %d%% optimized",
          metadata.getName(),
          metadata.getProgress().getProgressPercent()));
    } catch (InvalidProtocolBufferException e) {
      // The returned operation does not contain OptimizeRestoredDatabaseMetadata.
      System.err.println(e.getMessage());
    }
  }
}
 

@Test
public void cpuGpuConfig_emptyProvisionEnqueued() {

  // Arrange
  WorkerQueues queues = WorkerQueueConfigurations.gpuAndFallback();

  // Act
  Operation operation = Operation.newBuilder().build();
  SetMultimap<String, String> provisions = HashMultimap.create();
  boolean success = queues.enqueueOperation(operation, provisions);

  // Assert
  assertThat(success).isTrue();
  assertThat(queues.queueSize("GPU")).isEqualTo(0);
  assertThat(queues.queueSize("Other")).isEqualTo(1);
}
 

@Test
public void gpuConfig_validProvisionsEnqueued() {

  // Arrange
  WorkerQueues queues = WorkerQueueConfigurations.gpu();

  // Act
  Operation operation = Operation.newBuilder().build();
  SetMultimap<String, String> provisions = HashMultimap.create();
  provisions.put("foo", "bar");
  provisions.put("gpu_required", "only key matters");
  provisions.put("another", "provision");
  boolean success = queues.enqueueOperation(operation, provisions);

  // Assert
  assertThat(success).isTrue();
  assertThat(queues.queueSize("GPU")).isEqualTo(1);
}
 

public RedisShardBackplane(
    RedisShardBackplaneConfig config,
    String source,
    Function<Operation, Operation> onPublish,
    Function<Operation, Operation> onComplete,
    Predicate<Operation> isPrequeued,
    Predicate<Operation> isDispatched,
    Supplier<JedisCluster> jedisClusterFactory) {
  this.config = config;
  this.source = source;
  this.onPublish = onPublish;
  this.onComplete = onComplete;
  this.isPrequeued = isPrequeued;
  this.isDispatched = isDispatched;
  this.jedisClusterFactory = jedisClusterFactory;
}
 

private Builder(
    ExecuteResponse.Builder executeResponse,
    Operation operation,
    Poller poller,
    Path execDir,
    Action action,
    Command command,
    QueueEntry queueEntry) {
  this.executeResponse = executeResponse;
  this.operation = operation;
  this.poller = poller;
  this.execDir = execDir;
  this.action = action;
  this.command = command;
  this.queueEntry = queueEntry;
}
 

@Test
public void gpuConfig_unvalidProvisionNotEnqueued() {

  // Arrange
  WorkerQueues queues = WorkerQueueConfigurations.gpu();

  // Act
  Operation operation = Operation.newBuilder().build();
  SetMultimap<String, String> provisions = HashMultimap.create();
  provisions.put("invalid_key", "invalid_value");
  boolean success = queues.enqueueOperation(operation, provisions);

  // Assert
  assertThat(success).isFalse();
  assertThat(queues.queueSize("GPU")).isEqualTo(0);
}
 

static void listBackupOperations(InstanceAdminClient instanceAdminClient, DatabaseId databaseId) {
  Instance instance = instanceAdminClient.getInstance(databaseId.getInstanceId().getInstance());
  // Get create backup operations for the sample database.
  String filter =
      String.format(
          "(metadata.database:%s) AND "
              + "(metadata.@type:type.googleapis.com/"
              + "google.spanner.admin.database.v1.CreateBackupMetadata)",
          databaseId.getName());
  Page<Operation> operations = instance.listBackupOperations(Options.filter(filter));
  for (Operation op : operations.iterateAll()) {
    try {
      CreateBackupMetadata metadata = op.getMetadata().unpack(CreateBackupMetadata.class);
      System.out.println(
          String.format(
              "Backup %s on database %s pending: %d%% complete",
              metadata.getName(),
              metadata.getDatabase(),
              metadata.getProgress().getProgressPercent()));
    } catch (InvalidProtocolBufferException e) {
      // The returned operation does not contain CreateBackupMetadata.
      System.err.println(e.getMessage());
    }
  }
}
 

@Test
public void cpuGpuConfig_validGpuProvisionEnqueued() {

  // Arrange
  WorkerQueues queues = WorkerQueueConfigurations.gpuAndFallback();

  // Act
  Operation operation = Operation.newBuilder().build();
  SetMultimap<String, String> provisions = HashMultimap.create();
  provisions.put("gpu_required", "only key matters");
  boolean success = queues.enqueueOperation(operation, provisions);

  // Assert
  assertThat(success).isTrue();
  assertThat(queues.queueSize("GPU")).isEqualTo(1);
  assertThat(queues.queueSize("Other")).isEqualTo(0);
}
 

protected static ExecuteOperationMetadata expectExecuteOperationMetadata(Operation operation) {
  String name = operation.getName();
  Any metadata = operation.getMetadata();
  QueuedOperationMetadata queuedOperationMetadata = maybeQueuedOperationMetadata(name, metadata);
  if (queuedOperationMetadata != null) {
    return queuedOperationMetadata.getExecuteOperationMetadata();
  }
  ExecutingOperationMetadata executingOperationMetadata =
      maybeExecutingOperationMetadata(name, metadata);
  if (executingOperationMetadata != null) {
    return executingOperationMetadata.getExecuteOperationMetadata();
  }
  CompletedOperationMetadata completedOperationMetadata =
      maybeCompletedOperationMetadata(name, metadata);
  if (completedOperationMetadata != null) {
    return completedOperationMetadata.getExecuteOperationMetadata();
  }
  try {
    return operation.getMetadata().unpack(ExecuteOperationMetadata.class);
  } catch (InvalidProtocolBufferException e) {
    logger.log(
        Level.SEVERE, format("invalid execute operation metadata %s", operation.getName()), e);
  }
  return null;
}
 

@Test
public void cpuGpuConfig_invalidProvisionEnqueued() {

  // Arrange
  WorkerQueues queues = WorkerQueueConfigurations.gpuAndFallback();

  // Act
  Operation operation = Operation.newBuilder().build();
  SetMultimap<String, String> provisions = HashMultimap.create();
  provisions.put("invalid_key", "invalid_value");
  boolean success = queues.enqueueOperation(operation, provisions);

  // Assert
  assertThat(success).isTrue();
  assertThat(queues.queueSize("GPU")).isEqualTo(0);
  assertThat(queues.queueSize("Other")).isEqualTo(1);
}
 

private Answer<Void> answerWith(@Nullable Operation op, Status status) {
  return invocationOnMock -> {
    @SuppressWarnings("unchecked")
    StreamObserver<Operation> responseObserver =
        (StreamObserver<Operation>) invocationOnMock.getArguments()[1];
    if (op != null) {
      responseObserver.onNext(op);
    }
    if (status.isOk()) {
      responseObserver.onCompleted();
    } else {
      responseObserver.onError(status.asRuntimeException());
    }
    return null;
  };
}
 

protected Action expectAction(Operation operation) throws InterruptedException {
  try {
    ExecuteOperationMetadata metadata = expectExecuteOperationMetadata(operation);
    if (metadata == null) {
      return null;
    }
    ByteString actionBlob = getBlob(metadata.getActionDigest());
    if (actionBlob != null) {
      return Action.parseFrom(actionBlob);
    }
  } catch (IOException e) {
    Status status = Status.fromThrowable(e);
    if (status.getCode() != io.grpc.Status.Code.NOT_FOUND) {
      logger.log(Level.SEVERE, "error retrieving action", e);
    }
  }
  return null;
}
 

public void observe(Operation operation) {
  try {
    watcher.observe(operation);
    if (operation == null || operation.getDone()) {
      set(null);
    }
  } catch (Throwable t) {
    setException(t);
  }
}
 

public static Operation parseOperationJson(String operationJson) {
  if (operationJson == null) {
    return null;
  }
  try {
    Operation.Builder operationBuilder = Operation.newBuilder();
    operationParser.merge(operationJson, operationBuilder);
    return operationBuilder.build();
  } catch (InvalidProtocolBufferException e) {
    logger.log(Level.SEVERE, "error parsing operation from " + operationJson, e);
    return null;
  }
}
 

@Override
public void promoteCampaignDraft(
    PromoteCampaignDraftRequest request, StreamObserver<Operation> responseObserver) {
  Object response = responses.remove();
  if (response instanceof Operation) {
    requests.add(request);
    responseObserver.onNext((Operation) response);
    responseObserver.onCompleted();
  } else if (response instanceof Exception) {
    responseObserver.onError((Exception) response);
  } else {
    responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
  }
}
 

@Test
public void canceledOperationIsNoLongerOutstanding() throws IOException, InterruptedException {
  Operation operation = executeAction(createSimpleAction());

  // should appear in outstanding list
  ListOperationsRequest listRequest =
      ListOperationsRequest.newBuilder()
          .setName(INSTANCE_NAME + "/operations")
          .setPageSize(1024)
          .build();

  OperationsGrpc.OperationsBlockingStub operationsStub =
      OperationsGrpc.newBlockingStub(inProcessChannel);

  ListOperationsResponse listResponse = operationsStub.listOperations(listRequest);

  assertThat(Iterables.transform(listResponse.getOperationsList(), o -> o.getName()))
      .containsExactly(operation.getName());

  CancelOperationRequest cancelRequest =
      CancelOperationRequest.newBuilder().setName(operation.getName()).build();

  operationsStub.cancelOperation(cancelRequest);

  // should now be gone
  listResponse = operationsStub.listOperations(listRequest);

  assertThat(listResponse.getOperationsList()).isEmpty();
}
 

@Override
public void deleteOperation(String name) {
  synchronized (operationLock(name)) {
    Operation deletedOperation = completedOperations.remove(name);
    if (deletedOperation == null && outstandingOperations.contains(name)) {
      throw new IllegalStateException();
    }
  }
}
 

@Override
public void promoteCampaignExperiment(
    PromoteCampaignExperimentRequest request, StreamObserver<Operation> responseObserver) {
  Object response = responses.remove();
  if (response instanceof Operation) {
    requests.add(request);
    responseObserver.onNext((Operation) response);
    responseObserver.onCompleted();
  } else if (response instanceof Exception) {
    responseObserver.onError((Exception) response);
  } else {
    responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
  }
}
 

@Test
@SuppressWarnings("all")
public void createCampaignExperimentTest2() throws Exception {
  Empty expectedResponse = Empty.newBuilder().build();
  Operation resultOperation =
      Operation.newBuilder()
          .setName("createCampaignExperimentTest2")
          .setDone(true)
          .setResponse(Any.pack(expectedResponse))
          .build();
  mockCampaignExperimentService.addResponse(resultOperation);

  String customerId = "customerId-1772061412";
  CampaignExperiment campaignExperiment = CampaignExperiment.newBuilder().build();

  Empty actualResponse =
      client.createCampaignExperimentAsync(customerId, campaignExperiment).get();
  Assert.assertEquals(expectedResponse, actualResponse);

  List<AbstractMessage> actualRequests = mockCampaignExperimentService.getRequests();
  Assert.assertEquals(1, actualRequests.size());
  CreateCampaignExperimentRequest actualRequest =
      (CreateCampaignExperimentRequest) actualRequests.get(0);

  Assert.assertEquals(customerId, actualRequest.getCustomerId());
  Assert.assertEquals(campaignExperiment, actualRequest.getCampaignExperiment());
  Assert.assertTrue(
      channelProvider.isHeaderSent(
          ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
          GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}
 

private void withCancellation(
    ServerCallStreamObserver<Operation> serverCallStreamObserver, ListenableFuture<Void> future) {
  addCallback(
      future,
      new FutureCallback<Void>() {
        boolean isCancelled() {
          return serverCallStreamObserver.isCancelled() || Context.current().isCancelled();
        }

        @Override
        public void onSuccess(Void result) {
          if (!isCancelled()) {
            try {
              serverCallStreamObserver.onCompleted();
            } catch (Exception e) {
              onFailure(e);
            }
          }
        }

        @Override
        public void onFailure(Throwable t) {
          if (!isCancelled() && !(t instanceof CancellationException)) {
            logger.log(Level.WARNING, "error occurred during execution", t);
            serverCallStreamObserver.onError(Status.fromThrowable(t).asException());
          }
        }
      },
      Context.current().fixedContextExecutor(directExecutor()));
  serverCallStreamObserver.setOnCancelHandler(() -> future.cancel(false));
}
 

void assertResponse(ExecuteResponse executeResponse) throws Exception {
  ArgumentCaptor<Operation> operationCaptor = ArgumentCaptor.forClass(Operation.class);
  verify(mockBackplane, times(1)).putOperation(operationCaptor.capture(), eq(COMPLETED));
  Operation erroredOperation = operationCaptor.getValue();
  assertThat(erroredOperation.getDone()).isTrue();
  CompletedOperationMetadata completedMetadata =
      erroredOperation.getMetadata().unpack(CompletedOperationMetadata.class);
  assertThat(completedMetadata.getExecuteOperationMetadata().getStage()).isEqualTo(COMPLETED);
  assertThat(erroredOperation.getResponse().unpack(ExecuteResponse.class))
      .isEqualTo(executeResponse);
}
 

@Override
public void execute(ExecuteRequest request, StreamObserver<Operation> responseObserver) {
  final String opName = UUID.randomUUID().toString();
  ListenableFuture<ActionResult> future =
      executorService.submit(Context.current().wrap(() -> execute(request, opName)));
  operationsCache.put(opName, future);
  // Send the first operation.
  responseObserver.onNext(Operation.newBuilder().setName(opName).build());
  // When the operation completes, send the result.
  waitExecution(opName, future, responseObserver);
}
 

@Test
@SuppressWarnings("all")
public void createCampaignExperimentTest() throws Exception {
  Empty expectedResponse = Empty.newBuilder().build();
  Operation resultOperation =
      Operation.newBuilder()
          .setName("createCampaignExperimentTest")
          .setDone(true)
          .setResponse(Any.pack(expectedResponse))
          .build();
  mockCampaignExperimentService.addResponse(resultOperation);

  String customerId = "customerId-1772061412";
  CampaignExperiment campaignExperiment = CampaignExperiment.newBuilder().build();
  boolean validateOnly = false;

  Empty actualResponse =
      client.createCampaignExperimentAsync(customerId, campaignExperiment, validateOnly).get();
  Assert.assertEquals(expectedResponse, actualResponse);

  List<AbstractMessage> actualRequests = mockCampaignExperimentService.getRequests();
  Assert.assertEquals(1, actualRequests.size());
  CreateCampaignExperimentRequest actualRequest =
      (CreateCampaignExperimentRequest) actualRequests.get(0);

  Assert.assertEquals(customerId, actualRequest.getCustomerId());
  Assert.assertEquals(campaignExperiment, actualRequest.getCampaignExperiment());
  Assert.assertEquals(validateOnly, actualRequest.getValidateOnly());
  Assert.assertTrue(
      channelProvider.isHeaderSent(
          ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
          GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}
 

void publishReset(JedisCluster jedis, Operation operation) {
  Instant effectiveAt = Instant.now();
  Instant expiresAt = nextExpiresAt(effectiveAt);
  publish(
      jedis,
      operationChannel(operation.getName()),
      Instant.now(),
      OperationChange.newBuilder()
          .setReset(
              OperationChange.Reset.newBuilder()
                  .setExpiresAt(toTimestamp(expiresAt))
                  .setOperation(operation)
                  .build()));
}
 

protected void updateOperationWatchers(Operation operation) throws InterruptedException {
  if (operation.getDone()) {
    synchronized (operationLock(operation.getName())) {
      completedOperations.put(operation.getName(), operation);
      outstandingOperations.remove(operation.getName());
    }
  } else {
    outstandingOperations.put(operation.getName(), operation);
  }
}
 

protected void expireOperation(Operation operation) throws InterruptedException {
  ActionResult actionResult =
      ActionResult.newBuilder()
          .setExitCode(-1)
          .setStderrRaw(
              ByteString.copyFromUtf8(
                  "[BUILDFARM]: Action timed out with no response from worker"))
          .build();
  ExecuteResponse executeResponse =
      ExecuteResponse.newBuilder()
          .setResult(actionResult)
          .setStatus(
              com.google.rpc.Status.newBuilder().setCode(Code.DEADLINE_EXCEEDED.getNumber()))
          .build();
  ExecuteOperationMetadata metadata = expectExecuteOperationMetadata(operation);
  if (metadata == null) {
    throw new IllegalStateException(
        "Operation " + operation.getName() + " did not contain valid metadata");
  }
  metadata = metadata.toBuilder().setStage(ExecutionStage.Value.COMPLETED).build();
  putOperation(
      operation
          .toBuilder()
          .setDone(true)
          .setMetadata(Any.pack(metadata))
          .setResponse(Any.pack(executeResponse))
          .build());
}
 

@Override
protected Operation createOperation(ActionKey actionKey) {
  String name = createOperationName(UUID.randomUUID().toString());

  ExecuteOperationMetadata metadata =
      ExecuteOperationMetadata.newBuilder().setActionDigest(actionKey.getDigest()).build();

  return Operation.newBuilder()
      .setName(name)
      .setDone(false)
      .setMetadata(Any.pack(metadata))
      .build();
}
 

@Override
public void runMutateJob(
    RunMutateJobRequest request, StreamObserver<Operation> responseObserver) {
  Object response = responses.remove();
  if (response instanceof Operation) {
    requests.add(request);
    responseObserver.onNext((Operation) response);
    responseObserver.onCompleted();
  } else if (response instanceof Exception) {
    responseObserver.onError((Exception) response);
  } else {
    responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
  }
}
 

@Test
@SuppressWarnings("all")
public void runMutateJobTest() throws Exception {
  Empty expectedResponse = Empty.newBuilder().build();
  Operation resultOperation =
      Operation.newBuilder()
          .setName("runMutateJobTest")
          .setDone(true)
          .setResponse(Any.pack(expectedResponse))
          .build();
  mockMutateJobService.addResponse(resultOperation);

  String formattedResourceName =
      MutateJobServiceClient.formatMutateJobName("[CUSTOMER]", "[MUTATE_JOB]");

  Empty actualResponse = client.runMutateJobAsync(formattedResourceName).get();
  Assert.assertEquals(expectedResponse, actualResponse);

  List<AbstractMessage> actualRequests = mockMutateJobService.getRequests();
  Assert.assertEquals(1, actualRequests.size());
  RunMutateJobRequest actualRequest = (RunMutateJobRequest) actualRequests.get(0);

  Assert.assertEquals(formattedResourceName, actualRequest.getResourceName());
  Assert.assertTrue(
      channelProvider.isHeaderSent(
          ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
          GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
}