Dynamically configure SemaphoreBackPressureHandler with BackPressureLimiter (#1251)

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/AbstractContainerOptions.java

@@ -20,6 +20,7 @@
 import io.awspring.cloud.sqs.support.converter.MessagingMessageConverter;
 import io.awspring.cloud.sqs.support.converter.SqsMessagingMessageConverter;
 import java.time.Duration;
+import java.util.function.Supplier;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.lang.Nullable;
 import org.springframework.retry.backoff.BackOffPolicy;
@@ -47,12 +48,16 @@ public abstract class AbstractContainerOptions<O extends ContainerOptions<O, B>,
 
 	private final Duration maxDelayBetweenPolls;
 
+	private final Duration standbyLimitPollingInterval;
+
 	private final Duration listenerShutdownTimeout;
 
 	private final Duration acknowledgementShutdownTimeout;
 
 	private final BackPressureMode backPressureMode;
 
+	private final Supplier<BackPressureHandler> backPressureHandlerSupplier;
+
 	private final ListenerMode listenerMode;
 
 	private final MessagingMessageConverter<?> messageConverter;
@@ -80,10 +85,12 @@ protected AbstractContainerOptions(Builder<?, ?> builder) {
 		this.autoStartup = builder.autoStartup;
 		this.pollTimeout = builder.pollTimeout;
 		this.pollBackOffPolicy = builder.pollBackOffPolicy;
+		this.standbyLimitPollingInterval = builder.standbyLimitPollingInterval;
 		this.maxDelayBetweenPolls = builder.maxDelayBetweenPolls;
 		this.listenerShutdownTimeout = builder.listenerShutdownTimeout;
 		this.acknowledgementShutdownTimeout = builder.acknowledgementShutdownTimeout;
 		this.backPressureMode = builder.backPressureMode;
+		this.backPressureHandlerSupplier = builder.backPressureHandlerSupplier;
 		this.listenerMode = builder.listenerMode;
 		this.messageConverter = builder.messageConverter;
 		this.acknowledgementMode = builder.acknowledgementMode;
@@ -122,6 +129,11 @@ public BackOffPolicy getPollBackOffPolicy() {
 		return this.pollBackOffPolicy;
 	}
 
+	@Override
+	public Duration getStandbyLimitPollingInterval() {
+		return this.standbyLimitPollingInterval;
+	}
+
 	@Override
 	public Duration getMaxDelayBetweenPolls() {
 		return this.maxDelayBetweenPolls;
@@ -154,6 +166,11 @@ public BackPressureMode getBackPressureMode() {
 		return this.backPressureMode;
 	}
 
+	@Override
+	public Supplier<BackPressureHandler> getBackPressureHandlerSupplier() {
+		return this.backPressureHandlerSupplier;
+	}
+
 	@Override
 	public ListenerMode getListenerMode() {
 		return this.listenerMode;
@@ -206,6 +223,8 @@ protected abstract static class Builder<B extends ContainerOptionsBuilder<B, O>,
 
 		private static final BackOffPolicy DEFAULT_POLL_BACK_OFF_POLICY = buildDefaultBackOffPolicy();
 
+		private static final Duration DEFAULT_STANDBY_LIMIT_POLLING_INTERVAL = Duration.ofMillis(100);
+
 		private static final Duration DEFAULT_SEMAPHORE_TIMEOUT = Duration.ofSeconds(10);
 
 		private static final Duration DEFAULT_LISTENER_SHUTDOWN_TIMEOUT = Duration.ofSeconds(20);
@@ -214,6 +233,8 @@ protected abstract static class Builder<B extends ContainerOptionsBuilder<B, O>,
 
 		private static final BackPressureMode DEFAULT_THROUGHPUT_CONFIGURATION = BackPressureMode.AUTO;
 
+		private static final Supplier<BackPressureHandler> DEFAULT_BACKPRESSURE_LIMITER = null;
+
 		private static final ListenerMode DEFAULT_MESSAGE_DELIVERY_STRATEGY = ListenerMode.SINGLE_MESSAGE;
 
 		private static final MessagingMessageConverter<?> DEFAULT_MESSAGE_CONVERTER = new SqsMessagingMessageConverter();
@@ -230,10 +251,14 @@ protected abstract static class Builder<B extends ContainerOptionsBuilder<B, O>,
 
 		private BackOffPolicy pollBackOffPolicy = DEFAULT_POLL_BACK_OFF_POLICY;
 
+		private Duration standbyLimitPollingInterval = DEFAULT_STANDBY_LIMIT_POLLING_INTERVAL;
+
 		private Duration maxDelayBetweenPolls = DEFAULT_SEMAPHORE_TIMEOUT;
 
 		private BackPressureMode backPressureMode = DEFAULT_THROUGHPUT_CONFIGURATION;
 
+		private Supplier<BackPressureHandler> backPressureHandlerSupplier = DEFAULT_BACKPRESSURE_LIMITER;
+
 		private Duration listenerShutdownTimeout = DEFAULT_LISTENER_SHUTDOWN_TIMEOUT;
 
 		private Duration acknowledgementShutdownTimeout = DEFAULT_ACKNOWLEDGEMENT_SHUTDOWN_TIMEOUT;
@@ -272,6 +297,7 @@ protected Builder(AbstractContainerOptions<?, ?> options) {
 			this.listenerShutdownTimeout = options.listenerShutdownTimeout;
 			this.acknowledgementShutdownTimeout = options.acknowledgementShutdownTimeout;
 			this.backPressureMode = options.backPressureMode;
+			this.backPressureHandlerSupplier = options.backPressureHandlerSupplier;
 			this.listenerMode = options.listenerMode;
 			this.messageConverter = options.messageConverter;
 			this.acknowledgementMode = options.acknowledgementMode;
@@ -315,6 +341,13 @@ public B pollBackOffPolicy(BackOffPolicy pollBackOffPolicy) {
 			return self();
 		}
 
+		@Override
+		public B standbyLimitPollingInterval(Duration standbyLimitPollingInterval) {
+			Assert.notNull(standbyLimitPollingInterval, "standbyLimitPollingInterval cannot be null");
+			this.standbyLimitPollingInterval = standbyLimitPollingInterval;
+			return self();
+		}
+
 		@Override
 		public B maxDelayBetweenPolls(Duration maxDelayBetweenPolls) {
 			Assert.notNull(maxDelayBetweenPolls, "semaphoreAcquireTimeout cannot be null");
@@ -364,6 +397,12 @@ public B backPressureMode(BackPressureMode backPressureMode) {
 			return self();
 		}
 
+		@Override
+		public B backPressureHandlerSupplier(Supplier<BackPressureHandler> backPressureHandlerSupplier) {
+			this.backPressureHandlerSupplier = backPressureHandlerSupplier;
+			return self();
+		}
+
 		@Override
 		public B acknowledgementInterval(Duration acknowledgementInterval) {
 			Assert.notNull(acknowledgementInterval, "acknowledgementInterval cannot be null");

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/AbstractPipelineMessageListenerContainer.java

@@ -225,6 +225,10 @@ private TaskExecutor validateCustomExecutor(TaskExecutor taskExecutor) {
 	}
 
 	protected BackPressureHandler createBackPressureHandler() {
+		O containerOptions = getContainerOptions();
+		if (containerOptions.getBackPressureHandlerSupplier() != null) {
+			return containerOptions.getBackPressureHandlerSupplier().get();
+		}
 		return SemaphoreBackPressureHandler.builder().batchSize(getContainerOptions().getMaxMessagesPerPoll())
 				.totalPermits(getContainerOptions().getMaxConcurrentMessages())
 				.acquireTimeout(getContainerOptions().getMaxDelayBetweenPolls())

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/BackPressureHandler.java

@@ -29,20 +29,48 @@
 public interface BackPressureHandler {
 
 	/**
-	 * Request a number of permits. Each obtained permit allows the
+	 * Requests a number of permits. Each obtained permit allows the
 	 * {@link io.awspring.cloud.sqs.listener.source.MessageSource} to retrieve one message.
 	 * @param amount the amount of permits to request.
 	 * @return the amount of permits obtained.
 	 * @throws InterruptedException if the Thread is interrupted while waiting for permits.
 	 */
 	int request(int amount) throws InterruptedException;
 
+	/**
+	 * Releases the specified amount of permits for processed messages. Each message that has been processed should
+	 * release one permit, whether processing was successful or not.
+	 * <p>
+	 * This method can is called in the following use cases:
+	 * <ul>
+	 * <li>{@link ReleaseReason#LIMITED}: permits were not used because another BackPressureHandler has a lower permits
+	 * limit and the difference in permits needs to be returned.</li>
+	 * <li>{@link ReleaseReason#NONE_FETCHED}: none of the permits were actually used because no messages were retrieved
+	 * from SQS. Permits need to be returned.</li>
+	 * <li>{@link ReleaseReason#PARTIAL_FETCH}: some of the permits were used (some messages were retrieved from SQS).
+	 * The unused ones need to be returned. The amount to be returned might be {@literal 0}, in which case it means all
+	 * the permits will be used as the same number of messages were fetched from SQS.</li>
+	 * <li>{@link ReleaseReason#PROCESSED}: a message processing finished, successfully or not.</li>
+	 * </ul>
+	 * @param amount the amount of permits to release.
+	 * @param reason the reason why the permits were released.
+	 */
+	default void release(int amount, ReleaseReason reason) {
+		release(amount);
+	}
+
 	/**
 	 * Release the specified amount of permits. Each message that has been processed should release one permit, whether
 	 * processing was successful or not.
 	 * @param amount the amount of permits to release.
+	 *
+	 * @deprecated This method is deprecated and will not be called by the Spring Cloud AWS SQS listener anymore.
+	 * Implement {@link #release(int, ReleaseReason)} instead.
 	 */
-	void release(int amount);
+	@Deprecated
+	default void release(int amount) {
+		release(amount, ReleaseReason.PROCESSED);
+	}
 
 	/**
 	 * Attempts to acquire all permits up to the specified timeout. If successful, means all permits were returned and
@@ -52,4 +80,24 @@ public interface BackPressureHandler {
 	 */
 	boolean drain(Duration timeout);
 
+	enum ReleaseReason {
+		/**
+		 * Permits were not used because another BackPressureHandler has a lower permits limit and the difference need
+		 * to be aligned across all handlers.
+		 */
+		LIMITED,
+		/**
+		 * No messages were retrieved from SQS, so all permits need to be returned.
+		 */
+		NONE_FETCHED,
+		/**
+		 * Some messages were fetched from SQS. Unused permits need to be returned.
+		 */
+		PARTIAL_FETCH,
+		/**
+		 * The processing of one or more messages finished, successfully or not.
+		 */
+		PROCESSED;
+	}
+
 }

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/BatchAwareBackPressureHandler.java

@@ -35,13 +35,24 @@ public interface BatchAwareBackPressureHandler extends BackPressureHandler {
 	 * Release a batch of permits. This has the semantics of letting the {@link BackPressureHandler} know that all
 	 * permits from a batch are being released, in opposition to {@link #release(int)} in which any number of permits
 	 * can be specified.
+	 *
+	 * @deprecated This method is deprecated and will not be called by the Spring Cloud AWS SQS listener anymore.
+	 * Implement {@link BackPressureHandler#release(int, ReleaseReason)} instead.
 	 */
-	void releaseBatch();
+	@Deprecated
+	default void releaseBatch() {
+		release(getBatchSize(), ReleaseReason.NONE_FETCHED);
+	}
 
 	/**
 	 * Return the configured batch size for this handler.
 	 * @return the batch size.
+	 *
+	 * @deprecated This method is deprecated and will not be used by the Spring Cloud AWS SQS listener anymore.
 	 */
-	int getBatchSize();
+	@Deprecated
+	default int getBatchSize() {
+		return 0;
+	}
 
 }

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/CompositeBackPressureHandler.java

@@ -0,0 +1,149 @@
+/*
+ * Copyright 2013-2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.awspring.cloud.sqs.listener;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.ReentrantLock;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CompositeBackPressureHandler implements BatchAwareBackPressureHandler, IdentifiableContainerComponent {
+
+	private static final Logger logger = LoggerFactory.getLogger(CompositeBackPressureHandler.class);
+
+	private final List<BackPressureHandler> backPressureHandlers;
+
+	private final int batchSize;
+
+	private final ReentrantLock noPermitsReturnedWaitLock = new ReentrantLock();
+
+	private final Condition permitsReleasedCondition = noPermitsReturnedWaitLock.newCondition();
+
+	private final Duration noPermitsReturnedWaitTimeout;
+
+	private String id;
+
+	public CompositeBackPressureHandler(List<BackPressureHandler> backPressureHandlers, int batchSize,
+			Duration waitTimeout) {
+		this.backPressureHandlers = backPressureHandlers;
+		this.batchSize = batchSize;
+		this.noPermitsReturnedWaitTimeout = waitTimeout;
+	}
+
+	@Override
+	public void setId(String id) {
+		this.id = id;
+		backPressureHandlers.stream().filter(IdentifiableContainerComponent.class::isInstance)
+				.map(IdentifiableContainerComponent.class::cast)
+				.forEach(bph -> bph.setId(bph.getClass().getSimpleName() + "-" + id));
+	}
+
+	@Override
+	public String getId() {
+		return id;
+	}
+
+	@Override
+	public int requestBatch() throws InterruptedException {
+		return request(batchSize);
+	}
+
+	@Override
+	public int request(int amount) throws InterruptedException {
+		logger.debug("[{}] Requesting {} permits", this.id, amount);
+		int obtained = amount;
+		int[] obtainedPerBph = new int[backPressureHandlers.size()];
+		for (int i = 0; i < backPressureHandlers.size() && obtained > 0; i++) {
+			obtainedPerBph[i] = backPressureHandlers.get(i).request(obtained);
+			obtained = Math.min(obtained, obtainedPerBph[i]);
+		}
+		for (int i = 0; i < backPressureHandlers.size(); i++) {
+			int obtainedForBph = obtainedPerBph[i];
+			if (obtainedForBph > obtained) {
+				backPressureHandlers.get(i).release(obtainedForBph - obtained, ReleaseReason.LIMITED);
+			}
+		}
+		if (obtained == 0) {
+			waitForPermitsToBeReleased();
+		}
+		logger.debug("[{}] Obtained {} permits ({} requested)", this.id, obtained, amount);
+		return obtained;
+	}
+
+	@Override
+	public void release(int amount, ReleaseReason reason) {
+		logger.debug("[{}] Releasing {} permits ({})", this.id, amount, reason);
+		for (BackPressureHandler handler : backPressureHandlers) {
+			handler.release(amount, reason);
+		}
+		if (amount > 0) {
+			signalPermitsWereReleased();
+		}
+	}
+
+	/**
+	 * Waits for permits to be released up to {@link #noPermitsReturnedWaitTimeout}. If no permits were released within
+	 * the configured {@link #noPermitsReturnedWaitTimeout}, returns immediately. This allows {@link #request(int)} to
+	 * return {@code 0} permits and will trigger another round of back-pressure handling.
+	 *
+	 * @throws InterruptedException if the Thread is interrupted while waiting for permits.
+	 */
+	@SuppressWarnings({ "java:S899" // we are not interested in the await return value here
+	})
+	private void waitForPermitsToBeReleased() throws InterruptedException {
+		noPermitsReturnedWaitLock.lock();
+		try {
+			logger.trace("[{}] No permits were obtained, waiting for a release up to {}", this.id,
+					noPermitsReturnedWaitTimeout);
+			permitsReleasedCondition.await(noPermitsReturnedWaitTimeout.toMillis(), TimeUnit.MILLISECONDS);
+		}
+		finally {
+			noPermitsReturnedWaitLock.unlock();
+		}
+	}
+
+	private void signalPermitsWereReleased() {
+		noPermitsReturnedWaitLock.lock();
+		try {
+			permitsReleasedCondition.signal();
+		}
+		finally {
+			noPermitsReturnedWaitLock.unlock();
+		}
+	}
+
+	@Override
+	public boolean drain(Duration timeout) {
+		logger.debug("[{}] Draining back-pressure handlers initiated", this.id);
+		boolean result = true;
+		Instant start = Instant.now();
+		for (BackPressureHandler handler : backPressureHandlers) {
+			Duration remainingTimeout = maxDuration(timeout.minus(Duration.between(start, Instant.now())),
+					Duration.ZERO);
+			result &= handler.drain(remainingTimeout);
+		}
+		logger.debug("[{}] Draining back-pressure handlers completed", this.id);
+		return result;
+	}
+
+	private static Duration maxDuration(Duration first, Duration second) {
+		return first.compareTo(second) > 0 ? first : second;
+	}
+}