com.amazonaws.handlers.AsyncHandler Java Examples

private void unsubscribe(Subscription subscription) {
  if (PENDING_CONFIRMATION.equals(subscription.getSubscriptionArn())) {
    logger.warn("Could not remove subscription ({}) because it's pending for confirmation. It will be removed automatically "
        + "after 3 days.", subscription.getEndpoint());
    return;
  }
  SNS_CLIENT.unsubscribeAsync(subscription.getSubscriptionArn(), new AsyncHandler<UnsubscribeRequest, UnsubscribeResult>() {
    @Override
    public void onError(Exception e) {
      logger.error("Error un-subscribing endpoint {} from SNS: {}", subscription.getEndpoint(), e);
    }

    @Override
    public void onSuccess(UnsubscribeRequest request, UnsubscribeResult unsubscribeResult) {
      logger.debug("Endpoint {} has been successfully un-subscribed from SNS.", subscription.getEndpoint());
    }
  });
}
 

@Override
public void sendRawMessage(String message) {
  if (SNS_CLIENT == null) {
    logger.warn("The AdminMessage can not be sent as the MessageBroker is not ready. Message was: {}", message);
    return;
  }
  if (message.length() > MAX_MESSAGE_SIZE) {
    throw new RuntimeException("AdminMessage is larger than the MAX_MESSAGE_SIZE. Can not send it.");
  }
  //Send using SNS client
  SNS_CLIENT.publishAsync(TOPIC_ARN, message, new AsyncHandler<PublishRequest, PublishResult>() {
    @Override
    public void onError(Exception exception) {
      logger.error("Error sending message: {}", message, exception);
    }

    @Override
    public void onSuccess(PublishRequest request, PublishResult publishResult) {
      logger.debug("Message has been sent with following content: {}", message);
    }
  });
}
 

private static Future<CreateTableResult> createTableAsync(final AmazonDynamoDBAsync dbClient, final String tableName) {
  final String partitionKeyName = tableName + "Id";
  final CreateTableRequest createTableRequest = new CreateTableRequest()
    .withTableName(tableName).withKeySchema(new KeySchemaElement().withAttributeName(partitionKeyName).withKeyType(KeyType.HASH))
    .withAttributeDefinitions(new AttributeDefinition().withAttributeName(partitionKeyName).withAttributeType("S"))
    .withProvisionedThroughput(new ProvisionedThroughput().withReadCapacityUnits(10L).withWriteCapacityUnits(5L));

  return dbClient.createTableAsync(createTableRequest, new AsyncHandler<CreateTableRequest,CreateTableResult>() {
    @Override
    public void onError(final Exception exception) {
    }

    @Override
    public void onSuccess(final CreateTableRequest request, final CreateTableResult createTableResult) {
    }
  });
}
 

@SuppressWarnings("deprecation")
@Test
public void multipleRecordsEmitsTest() throws Exception {
    // Set up the buffer and do sanity checks
    buffer.clear();
    buffer.consumeRecord(ITEM1_INSERT, ITEM1_INSERT.getDynamodb().getSizeBytes().intValue(), ITEM1_INSERT.getDynamodb().getSequenceNumber());
    buffer.consumeRecord(ITEM2_INSERT, ITEM2_INSERT.getDynamodb().getSizeBytes().intValue(), ITEM2_INSERT.getDynamodb().getSequenceNumber());
    assertEquals(ITEM1_INSERT.getDynamodb().getSequenceNumber(), buffer.getFirstSequenceNumber());
    assertEquals(ITEM2_INSERT.getDynamodb().getSequenceNumber(), buffer.getLastSequenceNumber());
    List<Record> buffered = buffer.getRecords();
    assertEquals(2, buffered.size());
    assertTrue(buffered.contains(ITEM1_INSERT));
    assertTrue(buffered.contains(ITEM2_INSERT));

    // Emit record
    resetAll(DYNAMODB);
    DYNAMODB.putItemAsync(EasyMock.anyObject(PutItemRequest.class), anyObject(AsyncHandler.class));
    expectLastCall().andAnswer(SUCCESS_ANSWER).times(2);
    expectNew(AmazonDynamoDBAsyncClient.class, new Class<?>[] {AWSCredentialsProvider.class, ClientConfiguration.class, ExecutorService.class}, anyObject(AWSCredentialsProvider.class), anyObject(ClientConfiguration.class), anyObject(ExecutorService.class)).andReturn(DYNAMODB);
    DYNAMODB.setEndpoint(EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    replayAll(DYNAMODB);
    IEmitter<Record> instance = createEmitterInstance();
    assertTrue(instance.emit(new UnmodifiableBuffer<Record>(buffer)).isEmpty());
    verifyAll();
}
 

@SuppressWarnings("deprecation")
@Test
public void removeTest() throws Exception {
    // Set up the buffer and do sanity checks
    buffer.clear();
    buffer.consumeRecord(ITEM1_REMOVE, ITEM1_REMOVE.getDynamodb().getSizeBytes().intValue(), ITEM1_REMOVE.getDynamodb().getSequenceNumber());
    assertEquals(ITEM1_REMOVE.getDynamodb().getSequenceNumber(), buffer.getFirstSequenceNumber());
    assertEquals(ITEM1_REMOVE.getDynamodb().getSequenceNumber(), buffer.getLastSequenceNumber());
    List<Record> buffered = buffer.getRecords();
    assertEquals(1, buffered.size());
    assertTrue(buffered.contains(ITEM1_REMOVE));

    // Emit record
    resetAll(DYNAMODB);
    DYNAMODB.deleteItemAsync(anyObject(DeleteItemRequest.class), anyObject(AsyncHandler.class));
    expectLastCall().andAnswer(SUCCESS_ANSWER);
    expectNew(AmazonDynamoDBAsyncClient.class, new Class<?>[] {AWSCredentialsProvider.class, ClientConfiguration.class, ExecutorService.class}, anyObject(AWSCredentialsProvider.class), anyObject(ClientConfiguration.class), anyObject(ExecutorService.class)).andReturn(DYNAMODB);
    DYNAMODB.setEndpoint(EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    replayAll(DYNAMODB);
    IEmitter<Record> instance = createEmitterInstance();
    assertTrue(instance.emit(new UnmodifiableBuffer<Record>(buffer)).isEmpty());
    verifyAll();
}
 

@Bean
public AsyncHandler<PutRecordRequest, PutRecordResult> asyncHandler() {
	return new AsyncHandler<PutRecordRequest, PutRecordResult>() {

		@Override
		public void onError(Exception exception) {

		}

		@Override
		public void onSuccess(PutRecordRequest request, PutRecordResult putRecordsResult) {
			ProcessorConfiguration.this.resultMonoProcessor.onNext(putRecordsResult);
			ProcessorConfiguration.this.resultMonoProcessor.onComplete();
		}

	};
}
 

@SuppressWarnings("deprecation")
@Test
public void modifyTest() throws Exception {
    // Set up the buffer and do sanity checks
    buffer.clear();
    buffer.consumeRecord(ITEM1_MODIFY, ITEM1_MODIFY.getDynamodb().getSizeBytes().intValue(), ITEM1_MODIFY.getDynamodb().getSequenceNumber());
    assertEquals(ITEM1_MODIFY.getDynamodb().getSequenceNumber(), buffer.getFirstSequenceNumber());
    assertEquals(ITEM1_MODIFY.getDynamodb().getSequenceNumber(), buffer.getLastSequenceNumber());
    List<Record> buffered = buffer.getRecords();
    assertEquals(1, buffered.size());
    assertTrue(buffered.contains(ITEM1_MODIFY));

    // Emit record
    resetAll(DYNAMODB);
    DYNAMODB.putItemAsync(EasyMock.anyObject(PutItemRequest.class), anyObject(AsyncHandler.class));
    expectLastCall().andAnswer(SUCCESS_ANSWER);
    expectNew(AmazonDynamoDBAsyncClient.class, new Class<?>[] {AWSCredentialsProvider.class, ClientConfiguration.class, ExecutorService.class}, anyObject(AWSCredentialsProvider.class), anyObject(ClientConfiguration.class), anyObject(ExecutorService.class)).andReturn(DYNAMODB);
    DYNAMODB.setEndpoint(EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    replayAll(DYNAMODB);
    IEmitter<Record> instance = createEmitterInstance();
    assertTrue(instance.emit(new UnmodifiableBuffer<Record>(buffer)).isEmpty());
    verifyAll();
}
 

@SuppressWarnings("deprecation")
@Test
public void insertTest() throws Exception {
    // Set up the buffer and do sanity checks
    buffer.clear();
    buffer.consumeRecord(ITEM1_INSERT, ITEM1_INSERT.getDynamodb().getSizeBytes().intValue(), ITEM1_INSERT.getDynamodb().getSequenceNumber());
    assertEquals(ITEM1_INSERT.getDynamodb().getSequenceNumber(), buffer.getFirstSequenceNumber());
    assertEquals(ITEM1_INSERT.getDynamodb().getSequenceNumber(), buffer.getLastSequenceNumber());
    List<Record> buffered = buffer.getRecords();
    assertEquals(1, buffered.size());
    assertTrue(buffered.contains(ITEM1_INSERT));

    // Emit record
    resetAll(DYNAMODB);
    expectNew(AmazonDynamoDBAsyncClient.class, new Class<?>[] {AWSCredentialsProvider.class, ClientConfiguration.class, ExecutorService.class}, anyObject(AWSCredentialsProvider.class), anyObject(ClientConfiguration.class), anyObject(ExecutorService.class)).andReturn(DYNAMODB);
    DYNAMODB.putItemAsync(EasyMock.anyObject(PutItemRequest.class), anyObject(AsyncHandler.class));
    expectLastCall().andAnswer(SUCCESS_ANSWER);

    DYNAMODB.setEndpoint(EasyMock.anyString());
    EasyMock.expectLastCall().anyTimes();
    replayAll(DYNAMODB);
    IEmitter<Record> instance = createEmitterInstance();
    assertTrue(instance.emit(new UnmodifiableBuffer<Record>(buffer)).isEmpty());
    verifyAll();
}
 

Future<RES> someAsyncOperation(AsyncHandler<? super REQ, RES> handler) {
    futureTask = new FutureTask<RES>(() -> {
        handler.onSuccess(request, response);
        return response;
    });
    return futureTask;
}
 

private <T extends Serializable> void sendMessageToAllRegionsAsync(final T body ) throws IOException {
    if ( snsAsync == null ) {
        logger.error( "SNS client is null, perhaps it failed to initialize successfully" );
        return;
    }

    final String stringBody = toString( body );

    final String topicArn = getWriteTopicArn();

    if ( logger.isTraceEnabled() ) {
        logger.trace( "Publishing Message...{} to arn: {}", stringBody, topicArn );
    }

    PublishRequest publishRequest = new PublishRequest( topicArn, stringBody );

    snsAsync.publishAsync( publishRequest, new AsyncHandler<PublishRequest, PublishResult>() {
        @Override
        public void onError( Exception e ) {
            logger.error( "Error publishing message... {}", e );
            logger.error(FAILED_TO_SEND_MESSAGE, stringBody, topicArn, e);
        }


        @Override
        public void onSuccess( PublishRequest request, PublishResult result ) {
            if ( logger.isTraceEnabled() ) {
                logger.trace( "Successfully published... messageID=[{}],  arn=[{}]", result.getMessageId(),
                    request.getTopicArn() );
            }
        }
    } );
}
 

private void retryAsync(final Map<String, List<WriteRequest>> map) {
  for (final Entry<String, List<WriteRequest>> requests : map.entrySet()) {
    for (final WriteRequest r : requests.getValue()) {
      if (r.getPutRequest() != null) {

        /**
         * The code is pretty similar to retry. The only difference is retryAsync uses
         * putItemAsync instead of putItem
         */
        final PutItemRequest putRequest =
            new PutItemRequest(requests.getKey(), r.getPutRequest().getItem());
        final Future<PutItemResult> future =
            client.putItemAsync(putRequest, new AsyncHandler<PutItemRequest, PutItemResult>() {

              @Override
              public void onError(final Exception exception) {
                LOGGER.warn("Putitem Async failed in Dynamo");
                futureMap.remove(putRequest);
              }

              @Override
              public void onSuccess(final PutItemRequest request, final PutItemResult result) {
                if (futureMap.remove(request) == null) {
                  LOGGER.warn("Unable to delete PutItemRequest from futuresMap ");
                }

                return;
              }
            });

        futureMap.put(putRequest, future);
      }
    }
  }
}
 

public static  <REQUEST extends AmazonWebServiceRequest, RESPONSE> Mono<RESPONSE> toMono(
        REQUEST request,
        BiFunction<REQUEST, AsyncHandler<REQUEST, RESPONSE>, Future<RESPONSE>> callFun
) {
    Supplier<REQUEST> supplier = () -> request;
    return toMono(supplier, callFun);
}
 

public <REQ extends AmazonWebServiceRequest> AsyncHandler<REQ, RES> handler() {
    return new AsyncHandler<REQ, RES>() {
        @Override
        public void onError(Exception exception) {
            subscriber.onError(exception);
        }

        @Override
        public void onSuccess(REQ request, RES result) {
            subscriber.onSuccess(result);
        }
    };
}
 

public <REQ extends AmazonWebServiceRequest, RES> AsyncHandler<REQ, RES> handler(Action2<REQ, RES> onSuccessAction, Action1<Exception> onErrorAction) {
    return new AsyncHandler<REQ, RES>() {
        @Override
        public void onError(Exception exception) {
            onErrorAction.call(exception);
            subscriber.onError(exception);
        }

        @Override
        public void onSuccess(REQ request, RES result) {
            onSuccessAction.call(request, result);
            subscriber.onCompleted();
        }
    };
}
 

@Override
public Completable deleteScalingPolicy(String policyRefId, String jobId) {
    DeleteScalingPolicyRequest deleteScalingPolicyRequest = new DeleteScalingPolicyRequest();
    deleteScalingPolicyRequest.setResourceId(
            AWSAppAutoScalingUtil.buildGatewayResourceId(jobId,
                    awsAppScalingConfig.getAWSGatewayEndpointPrefix(),
                    awsAppScalingConfig.getRegion(),
                    awsAppScalingConfig.getAWSGatewayEndpointTargetStage()));
    deleteScalingPolicyRequest.setServiceNamespace(SERVICE_NAMESPACE);
    deleteScalingPolicyRequest.setScalableDimension(SCALABLE_DIMENSION);
    deleteScalingPolicyRequest.setPolicyName(buildScalingPolicyName(policyRefId, jobId));

    return RetryWrapper.wrapWithExponentialRetry(String.format("deleteScalingPolicy %s for job %s", policyRefId, jobId),
            Observable.create(emitter -> awsAppAutoScalingClientAsync.deleteScalingPolicyAsync(deleteScalingPolicyRequest, new AsyncHandler<DeleteScalingPolicyRequest, DeleteScalingPolicyResult>() {
                @Override
                public void onError(Exception exception) {
                    if (exception instanceof ObjectNotFoundException) {
                        logger.info("Scaling policy does not exist anymore for job/policyRefId {}/{}", jobId, policyRefId);
                        emitter.onCompleted();
                    } else {
                        logger.error("Delete scaling policy exception {} - {}", jobId, exception.getMessage());
                        awsAppAutoScalingMetrics.registerAwsDeletePolicyError(exception);
                        emitter.onError(AutoScalePolicyException.errorDeletingPolicy(policyRefId, exception.getMessage()));
                    }
                }

                @Override
                public void onSuccess(DeleteScalingPolicyRequest request, DeleteScalingPolicyResult deleteScalingPolicyResult) {
                    int httpStatusCode = deleteScalingPolicyResult.getSdkHttpMetadata().getHttpStatusCode();
                    logger.info("Deleted scaling policy for job/policyRefId {}/{}, status - {}", jobId, policyRefId, httpStatusCode);
                    awsAppAutoScalingMetrics.registerAwsDeletePolicySuccess();
                    emitter.onCompleted();

                }
            }), Emitter.BackpressureMode.NONE)).toCompletable();
}
 

@Override
public Observable<AutoScalableTarget> getScalableTargetsForJob(String jobId) {
    DescribeScalableTargetsRequest describeScalableTargetsRequest = new DescribeScalableTargetsRequest();
    describeScalableTargetsRequest.setServiceNamespace(SERVICE_NAMESPACE);
    describeScalableTargetsRequest.setScalableDimension(SCALABLE_DIMENSION);
    describeScalableTargetsRequest.setResourceIds(Collections.singletonList(
            AWSAppAutoScalingUtil.buildGatewayResourceId(jobId,
                    awsAppScalingConfig.getAWSGatewayEndpointPrefix(),
                    awsAppScalingConfig.getRegion(),
                    awsAppScalingConfig.getAWSGatewayEndpointTargetStage())));

    return RetryWrapper.wrapWithExponentialRetry(String.format("getScalableTargetsForJob for job %s", jobId),
            Observable.create(emitter -> awsAppAutoScalingClientAsync.describeScalableTargetsAsync(describeScalableTargetsRequest,
                    new AsyncHandler<DescribeScalableTargetsRequest, DescribeScalableTargetsResult>() {
                        @Override
                        public void onError(Exception exception) {
                            logger.error("Get scalable target exception for {} - {}", jobId, exception.getMessage());
                            awsAppAutoScalingMetrics.registerAwsGetTargetError(exception);
                            emitter.onError(exception);
                        }

                        @Override
                        public void onSuccess(DescribeScalableTargetsRequest request, DescribeScalableTargetsResult describeScalableTargetsResult) {
                            awsAppAutoScalingMetrics.registerAwsGetTargetSuccess();
                            List<ScalableTarget> scalableTargets = describeScalableTargetsResult.getScalableTargets();
                            scalableTargets.stream()
                                    .map(AWSAppAutoScalingUtil::toAutoScalableTarget)
                                    .forEach(emitter::onNext);
                            emitter.onCompleted();
                        }
                    }), Emitter.BackpressureMode.NONE));
}
 

private void loadSubscriptions(String nextToken, Handler<AsyncResult<List<Subscription>>> callback) {
  ListSubscriptionsByTopicRequest req = new ListSubscriptionsByTopicRequest()
      .withTopicArn(TOPIC_ARN);
  if (nextToken != null) {
    req.setNextToken(nextToken);
  }
  SNS_CLIENT.listSubscriptionsByTopicAsync(req,
      new AsyncHandler<ListSubscriptionsByTopicRequest, ListSubscriptionsByTopicResult>() {
        @Override
        public void onError(Exception e) {
          callback.handle(Future.failedFuture(e));
        }

        @Override
        public void onSuccess(ListSubscriptionsByTopicRequest request, ListSubscriptionsByTopicResult result) {
          List<Subscription> subscriptions = result.getSubscriptions()
              .stream()
              .filter(subscription -> SNS_HTTP_PROTOCOL.equals(subscription.getProtocol())
                  && subscription.getEndpoint().contains(AdminApi.ADMIN_MESSAGES_ENDPOINT))
              .collect(Collectors.toCollection(LinkedList::new));
          if (result.getNextToken() != null) {
            loadSubscriptions(result.getNextToken(), subscriptionsResult -> {
              if (subscriptionsResult.succeeded()) {
                subscriptions.addAll(subscriptionsResult.result());
                callback.handle(Future.succeededFuture(subscriptions));
              } else {
                callback.handle(subscriptionsResult);
              }
            });
          }
          else {
            callback.handle(Future.succeededFuture(subscriptions));
          }
        }
      });
}
 

/**
 * Invokes the remote lambda function and returns the decompressed response as bytes.
 */
@Override
protected void invoke(final Marker marker, byte[] bytes, boolean fireAndForget, final Handler<AsyncResult<byte[]>> callback) {
  final RemoteFunctionConfig remoteFunction = getConnectorConfig().remoteFunction;
  logger.debug(marker, "Invoking remote lambda function with id '{}' Event size is: {}", remoteFunction.id, bytes.length);

  InvokeRequest invokeReq = new InvokeRequest()
      .withFunctionName(((AWSLambda) remoteFunction).lambdaARN)
      .withPayload(ByteBuffer.wrap(bytes))
      .withInvocationType(fireAndForget ? InvocationType.Event : InvocationType.RequestResponse);

  asyncClient.invokeAsync(invokeReq, new AsyncHandler<InvokeRequest, InvokeResult>() {
    @Override
    public void onError(Exception exception) {
      if (callback == null) {
        logger.error(marker, "Error sending event to remote lambda function", exception);
      }
      else {
        callback.handle(Future.failedFuture(getWHttpException(marker, exception)));
      }
    }

    @Override
    public void onSuccess(InvokeRequest request, InvokeResult result) {
      byte[] responseBytes = new byte[result.getPayload().remaining()];
      result.getPayload().get(responseBytes);
      callback.handle(Future.succeededFuture(responseBytes));
    }
  });
}
 

void sendMetricData(List<MetricDatum> metricData) throws InterruptedException {
    PutMetricDataRequest putMetricDataRequest = new PutMetricDataRequest()
            .withNamespace(config.namespace())
            .withMetricData(metricData);
    CountDownLatch latch = new CountDownLatch(1);
    amazonCloudWatchAsync.putMetricDataAsync(putMetricDataRequest, new AsyncHandler<PutMetricDataRequest, PutMetricDataResult>() {
        @Override
        public void onError(Exception exception) {
            if (exception instanceof AbortedException) {
                logger.warn("sending metric data was aborted: {}", exception.getMessage());
            } else {
                logger.error("error sending metric data.", exception);
            }
            latch.countDown();
        }

        @Override
        public void onSuccess(PutMetricDataRequest request, PutMetricDataResult result) {
            logger.debug("published metric with namespace:{}", request.getNamespace());
            latch.countDown();
        }
    });
    try {
        @SuppressWarnings("deprecation")
        long readTimeoutMillis = config.readTimeout().toMillis();
        latch.await(readTimeoutMillis, TimeUnit.MILLISECONDS);
    } catch (InterruptedException e) {
        logger.warn("metrics push to cloudwatch took longer than expected");
        throw e;
    }
}
 

@Override
public Completable deleteAlarm(String policyRefId, String jobId) {
    DeleteAlarmsRequest deleteAlarmsRequest = new DeleteAlarmsRequest();
    deleteAlarmsRequest.setAlarmNames(Arrays.asList(buildCloudWatchName(policyRefId, jobId)));


    return wrapWithExponentialRetry(String.format("deleteAlarm in policy %s for job %s", policyRefId, jobId),
            Observable.create(emitter ->
                    awsCloudWatch.deleteAlarmsAsync(deleteAlarmsRequest, new AsyncHandler<DeleteAlarmsRequest, DeleteAlarmsResult>() {
                        @Override
                        public void onError(Exception exception) {
                            deleteErrorCounter.increment();
                            if (exception instanceof ResourceNotFoundException) {
                                emitter.onError(AutoScalePolicyException.unknownScalingPolicy(policyRefId, exception.getMessage()));
                            } else {
                                emitter.onError(AutoScalePolicyException.errorDeletingAlarm(policyRefId, exception.getMessage()));
                            }
                        }

                        @Override
                        public void onSuccess(DeleteAlarmsRequest request, DeleteAlarmsResult deleteAlarmsResult) {
                            int httpStatusCode = deleteAlarmsResult.getSdkHttpMetadata().getHttpStatusCode();
                            log.info("Deleted cloud watch alarm for job-id {}, status {}", jobId, httpStatusCode);
                            deleteAlarmCounter.increment();
                            emitter.onCompleted();
                        }
                    }), Emitter.BackpressureMode.NONE)).toCompletable();
}
 

@Override
public Completable deleteScalableTarget(String jobId) {
    DeregisterScalableTargetRequest deRegisterRequest = new DeregisterScalableTargetRequest();
    deRegisterRequest.setResourceId(
            AWSAppAutoScalingUtil.buildGatewayResourceId(jobId,
                    awsAppScalingConfig.getAWSGatewayEndpointPrefix(),
                    awsAppScalingConfig.getRegion(),
                    awsAppScalingConfig.getAWSGatewayEndpointTargetStage()));
    deRegisterRequest.setServiceNamespace(SERVICE_NAMESPACE);
    deRegisterRequest.setScalableDimension(SCALABLE_DIMENSION);

    return RetryWrapper.wrapWithExponentialRetry(String.format("deleteScalableTarget for job %s", jobId),
            Observable.create(emitter -> awsAppAutoScalingClientAsync.deregisterScalableTargetAsync(deRegisterRequest, new AsyncHandler<DeregisterScalableTargetRequest, DeregisterScalableTargetResult>() {
                @Override
                public void onError(Exception exception) {
                    if (exception instanceof ObjectNotFoundException) {
                        logger.info("Scalable target does not exist anymore for job {}", jobId);
                        emitter.onCompleted();
                    } else {
                        logger.error("Deregister scalable target exception {} - {}", jobId, exception.getMessage());
                        awsAppAutoScalingMetrics.registerAwsDeleteTargetError(exception);
                        emitter.onError(exception);
                    }
                }

                @Override
                public void onSuccess(DeregisterScalableTargetRequest request, DeregisterScalableTargetResult deregisterScalableTargetResult) {
                    int httpStatusCode = deregisterScalableTargetResult.getSdkHttpMetadata().getHttpStatusCode();
                    logger.info("De-registered scalable target for {}, status {}", jobId, httpStatusCode);
                    awsAppAutoScalingMetrics.registerAwsDeleteTargetSuccess();
                    emitter.onCompleted();
                }
            }), Emitter.BackpressureMode.NONE)).toCompletable();
}
 

@Override
public Completable createScalableTarget(String jobId, int minCapacity, int maxCapacity) {
    RegisterScalableTargetRequest registerScalableTargetRequest = new RegisterScalableTargetRequest();
    registerScalableTargetRequest.setMinCapacity(minCapacity);
    registerScalableTargetRequest.setMaxCapacity(maxCapacity);
    registerScalableTargetRequest.setResourceId(AWSAppAutoScalingUtil.buildGatewayResourceId(jobId,
            awsAppScalingConfig.getAWSGatewayEndpointPrefix(),
            awsAppScalingConfig.getRegion(),
            awsAppScalingConfig.getAWSGatewayEndpointTargetStage()));
    registerScalableTargetRequest.setServiceNamespace(SERVICE_NAMESPACE);
    registerScalableTargetRequest.setScalableDimension(SCALABLE_DIMENSION);
    logger.info("RegisterScalableTargetRequest {}", registerScalableTargetRequest);

    return RetryWrapper.wrapWithExponentialRetry(String.format("createScalableTarget for job %s", jobId),
            Observable.create(emitter -> awsAppAutoScalingClientAsync.registerScalableTargetAsync(registerScalableTargetRequest, new AsyncHandler<RegisterScalableTargetRequest, RegisterScalableTargetResult>() {
                @Override
                public void onError(Exception exception) {
                    logger.error("Register scalable target exception for {} - {}", jobId, exception.getMessage());
                    awsAppAutoScalingMetrics.registerAwsCreateTargetError(exception);
                    emitter.onError(exception);
                }

                @Override
                public void onSuccess(RegisterScalableTargetRequest request, RegisterScalableTargetResult registerScalableTargetResult) {
                    int httpStatusCode = registerScalableTargetResult.getSdkHttpMetadata().getHttpStatusCode();
                    logger.info("Registered scalable target for success {} - status {}", jobId, httpStatusCode);
                    awsAppAutoScalingMetrics.registerAwsCreateTargetSuccess();
                    emitter.onCompleted();
                }
            }), Emitter.BackpressureMode.NONE)).toCompletable();
}
 

private <REQUEST extends AmazonWebServiceRequest, RESPONSE> Observable<RESPONSE> toObservable(
        REQUEST request,
        BiFunction<REQUEST, AsyncHandler<REQUEST, RESPONSE>, Future<RESPONSE>> callFun
) {
    Supplier<REQUEST> supplier = () -> request;
    return toObservable(supplier, callFun);
}
 

@Test
void executeMessage_successfulExecution_shouldRemoveMessageFromQueue()
		throws Exception {
	// Arrange
	CountDownLatch countDownLatch = new CountDownLatch(1);
	SimpleMessageListenerContainer container = new SimpleMessageListenerContainer() {

		@Override
		protected void executeMessage(
				org.springframework.messaging.Message<String> stringMessage) {
			countDownLatch.countDown();
			super.executeMessage(stringMessage);
		}
	};

	AmazonSQSAsync sqs = mock(AmazonSQSAsync.class);
	container.setAmazonSqs(sqs);

	QueueMessageHandler messageHandler = new QueueMessageHandler();
	container.setMessageHandler(messageHandler);

	StaticApplicationContext applicationContext = new StaticApplicationContext();
	applicationContext.registerSingleton("testMessageListener",
			TestMessageListener.class);

	mockGetQueueUrl(sqs, "testQueue",
			"https://executeMessage_successfulExecution_shouldRemoveMessageFromQueue.amazonaws.com");
	mockGetQueueAttributesWithEmptyResult(sqs,
			"https://executeMessage_successfulExecution_shouldRemoveMessageFromQueue.amazonaws.com");

	messageHandler.setApplicationContext(applicationContext);
	messageHandler.afterPropertiesSet();
	container.afterPropertiesSet();

	mockReceiveMessage(sqs,
			"https://executeMessage_successfulExecution_shouldRemoveMessageFromQueue.amazonaws.com",
			"messageContent", "ReceiptHandle");

	// Act
	container.start();

	// Assert
	assertThat(countDownLatch.await(2L, TimeUnit.SECONDS)).isTrue();
	container.stop();
	verify(sqs, times(1)).deleteMessageAsync(eq(new DeleteMessageRequest(
			"https://executeMessage_successfulExecution_shouldRemoveMessageFromQueue.amazonaws.com",
			"ReceiptHandle")), any(AsyncHandler.class));
}
 

/**
 * Fire the async query On success, we check to see if we can fire any more queries We continue to
 * fire queries until the global max is reached or we have asynchronously fired all queries
 *
 * <p> Any waiting threads are signaled here
 */
private void makeAsyncQuery() {
  synchronized (monitorLock) {
    ++asyncRequestsInProgress;
    dynamoDBClient.queryAsync(lastRequest, new AsyncHandler<QueryRequest, QueryResult>() {

      /**
       * On Error, add a null and notify the thread waiting This makes sure that they are not
       * stuck waiting
       */
      @Override
      public void onError(final Exception exception) {
        LOGGER.error("Query async failed with Exception ", exception);
        synchronized (monitorLock) {
          --asyncRequestsInProgress;
          decTotalAsyncRequestsInProgress();
          asyncQueryResults.add(null);
          monitorLock.notify();
        }
      }

      /**
       * On Success, fire a new request if we can Notify the waiting thread with the result
       */
      @Override
      public void onSuccess(final QueryRequest request, final QueryResult result) {

        synchronized (monitorLock) {
          --asyncRequestsInProgress;
          decTotalAsyncRequestsInProgress();

          if ((result.getLastEvaluatedKey() != null) && !result.getLastEvaluatedKey().isEmpty()) {
            lastRequest.setExclusiveStartKey(result.getLastEvaluatedKey());
            checkAndAsyncQuery();
          } else {
            lastRequest = null;
          }

          asyncQueryResults.add(result);
          monitorLock.notify();
        }
      }
    });
  }
  return;
}
 

/**
 * Fire the async query On success, we check to see if we can fire any more queries We continue to
 * fire queries until the global max is reached or we have asynchronously fired all queries
 *
 * <p> Any waiting threads are signaled here
 */
private void makeAsyncScan() {
  synchronized (monitorLock) {
    ++asyncRequestsInProgress;
    dynamoDBClient.scanAsync(lastRequest, new AsyncHandler<ScanRequest, ScanResult>() {

      /**
       * On Error, add a null and notify the thread waiting This makes sure that they are not
       * stuck waiting
       */
      @Override
      public void onError(final Exception exception) {
        LOGGER.error("Query async failed with Exception ", exception);
        synchronized (monitorLock) {
          --asyncRequestsInProgress;
          decTotalAsyncRequestsInProgress();
          asyncScanResults.add(null);
          monitorLock.notify();
        }
      }

      /**
       * On Success, fire a new request if we can Notify the waiting thread with the result
       */
      @Override
      public void onSuccess(final ScanRequest request, final ScanResult result) {

        synchronized (monitorLock) {
          --asyncRequestsInProgress;
          decTotalAsyncRequestsInProgress();

          if ((result.getLastEvaluatedKey() != null) && !result.getLastEvaluatedKey().isEmpty()) {
            lastRequest.setExclusiveStartKey(result.getLastEvaluatedKey());
            checkAndAsyncScan();
          } else {
            lastRequest = null;
          }

          asyncScanResults.add(result);
          monitorLock.notify();
        }
      }
    });
  }
  return;
}
 

private void writeBatch(final boolean async) {
  final List<WriteRequest> batch;

  if (batchedItems.size() <= NUM_ITEMS) {
    batch = batchedItems;
  } else {
    batch = batchedItems.subList(0, NUM_ITEMS + 1);
  }
  final Map<String, List<WriteRequest>> writes = new HashMap<>();
  writes.put(tableName, new ArrayList<>(batch));
  if (async) {

    /**
     * To support asynchronous batch write a async handler is created Callbacks are provided for
     * success and error. As there might be unprocessed items on failure, they are retried
     * asynchronously Keep track of futures, so that they can be waited on during "flush"
     */
    final BatchWriteItemRequest batchRequest = new BatchWriteItemRequest(writes);
    final Future<BatchWriteItemResult> future =
        client.batchWriteItemAsync(
            batchRequest,
            new AsyncHandler<BatchWriteItemRequest, BatchWriteItemResult>() {

              @Override
              public void onError(final Exception exception) {
                LOGGER.warn(
                    "Unable to get response from Dynamo-Async Write " + exception.toString());
                futureMap.remove(batchRequest);
                return;
              }

              @Override
              public void onSuccess(
                  final BatchWriteItemRequest request,
                  final BatchWriteItemResult result) {
                retryAsync(result.getUnprocessedItems());
                if (futureMap.remove(request) == null) {
                  LOGGER.warn(" Unable to delete BatchWriteRequest from futuresMap ");
                }
              }
            });

    futureMap.put(batchRequest, future);
  } else {
    final BatchWriteItemResult response =
        client.batchWriteItem(new BatchWriteItemRequest(writes));
    retry(response.getUnprocessedItems());
  }

  batch.clear();
}
 

@Test
void receiveMessage_withMessageListenerMethodAndNeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion()
		throws Exception {
	// Arrange
	CountDownLatch countDownLatch = new CountDownLatch(1);
	SimpleMessageListenerContainer container = new SimpleMessageListenerContainer() {

		@Override
		protected void executeMessage(
				org.springframework.messaging.Message<String> stringMessage) {
			countDownLatch.countDown();
			super.executeMessage(stringMessage);
		}
	};

	AmazonSQSAsync sqs = mock(AmazonSQSAsync.class);
	container.setAmazonSqs(sqs);

	QueueMessageHandler messageHandler = new QueueMessageHandler();
	container.setMessageHandler(messageHandler);

	StaticApplicationContext applicationContext = new StaticApplicationContext();
	applicationContext.registerSingleton("testListener",
			TestMessageListenerWithManualDeletionPolicy.class);

	mockGetQueueUrl(sqs, "testQueue",
			"https://receiveMessage_withMessageListenerMethodAnd"
					+ "NeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion.amazonaws.com");
	mockGetQueueAttributesWithEmptyResult(sqs,
			"https://receiveMessage_withMessageListenerMethodAnd"
					+ "NeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion.amazonaws.com");

	messageHandler.setApplicationContext(applicationContext);
	messageHandler.afterPropertiesSet();
	container.afterPropertiesSet();

	mockReceiveMessage(sqs, "https://receiveMessage_withMessageListenerMethodAnd"
			+ "NeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion.amazonaws.com",
			"messageContent", "ReceiptHandle");

	// Act
	container.start();

	// Assert
	countDownLatch.await(1L, TimeUnit.SECONDS);
	verify(sqs, never()).deleteMessageAsync(eq(new DeleteMessageRequest(
			"https://receiveMessage_withMessageListenerMethodAnd"
					+ "NeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion.amazonaws.com",
			"ReceiptHandle")), any(AsyncHandler.class));
	TestMessageListenerWithManualDeletionPolicy testMessageListenerWithManualDeletionPolicy = applicationContext
			.getBean(TestMessageListenerWithManualDeletionPolicy.class);
	testMessageListenerWithManualDeletionPolicy.getCountDownLatch().await(1L,
			TimeUnit.SECONDS);
	testMessageListenerWithManualDeletionPolicy.acknowledge();
	verify(sqs, times(1)).deleteMessageAsync(eq(new DeleteMessageRequest(
			"https://receiveMessage_withMessageListenerMethodAnd"
					+ "NeverDeletionPolicy_waitsForAcknowledgmentBeforeDeletion.amazonaws.com",
			"ReceiptHandle")), any(AsyncHandler.class));
	container.stop();
}
 

@Test
void executeMessage_executionThrowsExceptionAndQueueHasRedrivePolicy_shouldNotRemoveMessageFromQueue()
		throws Exception {
	// Arrange
	CountDownLatch countDownLatch = new CountDownLatch(1);
	SimpleMessageListenerContainer container = new SimpleMessageListenerContainer() {

		@Override
		protected void executeMessage(
				org.springframework.messaging.Message<String> stringMessage) {
			countDownLatch.countDown();
			super.executeMessage(stringMessage);
		}
	};

	AmazonSQSAsync sqs = mock(AmazonSQSAsync.class);
	container.setAmazonSqs(sqs);

	QueueMessageHandler messageHandler = new QueueMessageHandler();
	container.setMessageHandler(messageHandler);

	StaticApplicationContext applicationContext = new StaticApplicationContext();
	Class clazz = TestMessageListenerThatThrowsAnExceptionWithAllExceptOnRedriveDeletionPolicy.class;
	applicationContext.registerSingleton("testMessageListener", clazz);

	mockGetQueueUrl(sqs, "testQueue",
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasRedrivePolicy_shouldNotRemoveMessageFromQueue.amazonaws.com");
	mockGetQueueAttributesWithRedrivePolicy(sqs,
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasRedrivePolicy_shouldNotRemoveMessageFromQueue.amazonaws.com");

	messageHandler.setApplicationContext(applicationContext);
	messageHandler.afterPropertiesSet();
	container.afterPropertiesSet();

	when(sqs.receiveMessage(new ReceiveMessageRequest(
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasRedrivePolicy_shouldNotRemoveMessageFromQueue.amazonaws.com")
							.withAttributeNames("All").withMaxNumberOfMessages(10)
							.withWaitTimeSeconds(20)
							.withMessageAttributeNames("All")))
									.thenReturn(
											new ReceiveMessageResult()
													.withMessages(new Message()
															.withBody(
																	"messageContent")
															.withReceiptHandle(
																	"ReceiptHandle")),
											new ReceiveMessageResult());

	// Act
	container.start();

	// Assert
	assertThat(countDownLatch.await(2L, TimeUnit.SECONDS)).isTrue();
	container.stop();
	verify(sqs, never()).deleteMessageAsync(eq(new DeleteMessageRequest(
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasRedrivePolicy_shouldNotRemoveMessageFromQueue.amazonaws.com",
			"ReceiptHandle")), any(AsyncHandler.class));
}
 

@Test
void executeMessage_executionThrowsExceptionAndQueueHasAllDeletionPolicy_shouldRemoveMessageFromQueue()
		throws Exception {
	// Arrange
	CountDownLatch countDownLatch = new CountDownLatch(1);
	SimpleMessageListenerContainer container = new SimpleMessageListenerContainer() {

		@Override
		protected void executeMessage(
				org.springframework.messaging.Message<String> stringMessage) {
			countDownLatch.countDown();
			super.executeMessage(stringMessage);
		}
	};

	AmazonSQSAsync sqs = mock(AmazonSQSAsync.class);
	container.setAmazonSqs(sqs);

	QueueMessageHandler messageHandler = new QueueMessageHandler();
	container.setMessageHandler(messageHandler);

	StaticApplicationContext applicationContext = new StaticApplicationContext();
	applicationContext.registerSingleton("testMessageListener",
			TestMessageListenerThatThrowsAnExceptionWithAllDeletionPolicy.class);

	mockGetQueueUrl(sqs, "testQueue",
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasAllDeletionPolicy_shouldRemoveMessageFromQueue.amazonaws.com");
	mockGetQueueAttributesWithEmptyResult(sqs,
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasAllDeletionPolicy_shouldRemoveMessageFromQueue.amazonaws.com");

	messageHandler.setApplicationContext(applicationContext);
	messageHandler.afterPropertiesSet();
	container.afterPropertiesSet();

	when(sqs.receiveMessage(new ReceiveMessageRequest(
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasAllDeletionPolicy_shouldRemoveMessageFromQueue.amazonaws.com")
							.withAttributeNames("All").withMaxNumberOfMessages(10)
							.withWaitTimeSeconds(20)
							.withMessageAttributeNames("All")))
									.thenReturn(
											new ReceiveMessageResult()
													.withMessages(new Message()
															.withBody(
																	"messageContent")
															.withReceiptHandle(
																	"ReceiptHandle")),
											new ReceiveMessageResult());

	// Act
	container.start();

	// Assert
	assertThat(countDownLatch.await(2L, TimeUnit.SECONDS)).isTrue();
	container.stop();
	verify(sqs, times(1)).deleteMessageAsync(eq(new DeleteMessageRequest(
			"https://executeMessage_executionThrowsExceptionAnd"
					+ "QueueHasAllDeletionPolicy_shouldRemoveMessageFromQueue.amazonaws.com",
			"ReceiptHandle")), any(AsyncHandler.class));
}