Support precise sub second timeout for BZPOP.

Author: christophstrobl

src/main/java/org/springframework/data/redis/connection/ReactiveZSetCommands.java

@@ -25,10 +25,10 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
-
 import org.springframework.data.domain.Range;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.redis.connection.ReactiveRedisConnection.CommandResponse;
@@ -163,8 +163,7 @@ public ZAddCommand incr() {
 		}
 
 		/**
-		 * Applies {@literal GT} mode. Constructs a new command
-		 * instance with all previously configured properties.
+		 * Applies {@literal GT} mode. Constructs a new command instance with all previously configured properties.
 		 *
 		 * @return a new {@link ZAddCommand} with {@literal incr} applied.
 		 * @since 2.5
@@ -174,8 +173,7 @@ public ZAddCommand gt() {
 		}
 
 		/**
-		 * Applies {@literal LT} mode. Constructs a new command
-		 * instance with all previously configured properties.
+		 * Applies {@literal LT} mode. Constructs a new command instance with all previously configured properties.
 		 *
 		 * @return a new {@link ZAddCommand} with {@literal incr} applied.
 		 * @since 2.5
@@ -206,7 +204,6 @@ public boolean isIncr() {
 		}
 
 		/**
-		 *
 		 * @return {@literal true} if {@literal GT} is set.
 		 * @since 2.5
 		 */
@@ -1337,15 +1334,18 @@ class BZPopCommand extends KeyCommand {
 
 		private final PopDirection direction;
 
-		private final Duration timeout;
+		private final @Nullable TimeUnit timeUnit;
+		private final @Nullable Long timeout;
 
 		private final long count;
 
-		private BZPopCommand(@Nullable ByteBuffer key, Duration timeout, long count, PopDirection direction) {
+		private BZPopCommand(@Nullable ByteBuffer key, @Nullable Long timeout, @Nullable TimeUnit timeUnit, long count,
+				PopDirection direction) {
 
 			super(key);
 			this.count = count;
 			this.timeout = timeout;
+			this.timeUnit = timeUnit;
 			this.direction = direction;
 		}
 
@@ -1355,7 +1355,7 @@ private BZPopCommand(@Nullable ByteBuffer key, Duration timeout, long count, Pop
 		 * @return a new {@link BZPopCommand} for min pop ({@literal ZPOPMIN}).
 		 */
 		public static BZPopCommand min() {
-			return new BZPopCommand(null, null, 0, PopDirection.MIN);
+			return new BZPopCommand(null, null, null, 0, PopDirection.MIN);
 		}
 
 		/**
@@ -1364,7 +1364,7 @@ public static BZPopCommand min() {
 		 * @return a new {@link BZPopCommand} for max pop ({@literal ZPOPMAX}).
 		 */
 		public static BZPopCommand max() {
-			return new BZPopCommand(null, null, 0, PopDirection.MAX);
+			return new BZPopCommand(null, null, null, 0, PopDirection.MAX);
 		}
 
 		/**
@@ -1377,7 +1377,7 @@ public BZPopCommand from(ByteBuffer key) {
 
 			Assert.notNull(key, "Key must not be null!");
 
-			return new BZPopCommand(key, timeout, count, direction);
+			return new BZPopCommand(key, timeout, timeUnit, count, direction);
 		}
 
 		/**
@@ -1387,7 +1387,7 @@ public BZPopCommand from(ByteBuffer key) {
 		 * @return a new {@link BZPopCommand} with {@literal value} applied.
 		 */
 		public BZPopCommand count(long count) {
-			return new BZPopCommand(getKey(), timeout, count, direction);
+			return new BZPopCommand(getKey(), timeout, timeUnit, count, direction);
 		}
 
 		/**
@@ -1400,7 +1400,21 @@ public BZPopCommand blockingFor(Duration timeout) {
 
 			Assert.notNull(timeout, "Timeout must not be null!");
 
-			return new BZPopCommand(getKey(), timeout, count, direction);
+			return blockingFor(timeout.toMillis(), TimeUnit.MILLISECONDS);
+		}
+
+		/**
+		 * Applies a {@link Duration timeout}. Constructs a new command instance with all previously configured properties.
+		 *
+		 * @param timeout value.
+		 * @param timeout must not be {@literal null}.
+		 * @return a new {@link BZPopCommand} with {@link Duration timeout} applied.
+		 */
+		public BZPopCommand blockingFor(long timeout, TimeUnit timeUnit) {
+
+			Assert.notNull(timeUnit, "TimeUnit must not be null!");
+
+			return new BZPopCommand(getKey(), timeout, timeUnit, count, direction);
 		}
 
 		/**
@@ -1410,10 +1424,16 @@ public PopDirection getDirection() {
 			return direction;
 		}
 
-		public Duration getTimeout() {
+		@Nullable
+		public Long getTimeout() {
 			return timeout;
 		}
 
+		@Nullable
+		public TimeUnit getTimeUnit() {
+			return timeUnit;
+		}
+
 		public long getCount() {
 			return count;
 		}
@@ -1664,7 +1684,7 @@ public static ZMScoreCommand scoreOf(ByteBuffer member) {
 		/**
 		 * Creates a new {@link ZMScoreCommand} given a {@link List members}.
 		 *
-		 * @param member must not be {@literal null}.
+		 * @param members must not be {@literal null}.
 		 * @return a new {@link ZMScoreCommand} for {@link List} of members.
 		 */
 		public static ZMScoreCommand scoreOf(Collection<ByteBuffer> members) {

src/main/java/org/springframework/data/redis/connection/RedisZSetCommands.java

@@ -997,8 +997,8 @@ default Long zCount(byte[] key, double min, double max) {
 	Set<Tuple> zPopMin(byte[] key, long count);
 
 	/**
-	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout
@@ -1034,8 +1034,8 @@ default Long zCount(byte[] key, double min, double max) {
 	Set<Tuple> zPopMax(byte[] key, long count);
 
 	/**
-	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout

src/main/java/org/springframework/data/redis/connection/lettuce/LettuceReactiveZSetCommands.java

@@ -25,8 +25,8 @@
 import reactor.core.publisher.Mono;
 
 import java.nio.ByteBuffer;
-import java.time.temporal.ChronoUnit;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import org.reactivestreams.Publisher;
 
@@ -390,7 +390,18 @@ public Flux<CommandResponse<BZPopCommand, Flux<Tuple>>> bZPop(Publisher<BZPopCom
 			Assert.notNull(command.getKey(), "Key must not be null!");
 			Assert.notNull(command.getTimeout(), "Timeout must not be null!");
 
-			long timeout = command.getTimeout().get(ChronoUnit.SECONDS);
+			if(command.getTimeUnit() == TimeUnit.MILLISECONDS) {
+
+				double timeout = preciseTimeout(command.getTimeout(), command.getTimeUnit());
+
+				Mono<ScoredValue<ByteBuffer>> result = (command.getDirection() == PopDirection.MIN
+						? cmd.bzpopmin(timeout, command.getKey())
+						: cmd.bzpopmax(timeout, command.getKey())).filter(Value::hasValue).map(Value::getValue);
+
+				return new CommandResponse<>(command, result.filter(Value::hasValue).map(this::toTuple).flux());
+			}
+
+			long timeout = command.getTimeUnit().toSeconds(command.getTimeout());
 
 			Mono<ScoredValue<ByteBuffer>> result = (command.getDirection() == PopDirection.MIN
 					? cmd.bzpopmin(timeout, command.getKey())
@@ -624,6 +635,10 @@ private Tuple toTuple(ByteBuffer value, double score) {
 		return new DefaultTuple(ByteUtils.getBytes(value), score);
 	}
 
+	static double preciseTimeout(long val, TimeUnit unit) {
+		return (double) unit.toMillis(val) / 1000.0D;
+	}
+
 	protected LettuceReactiveRedisConnection getConnection() {
 		return connection;
 	}

src/main/java/org/springframework/data/redis/connection/lettuce/LettuceZSetCommands.java

@@ -316,6 +316,13 @@ public Tuple bZPopMin(byte[] key, long timeout, TimeUnit unit) {
 		Assert.notNull(key, "Key must not be null!");
 		Assert.notNull(unit, "TimeUnit must not be null!");
 
+		if(TimeUnit.MILLISECONDS == unit) {
+
+			return connection.invoke(connection.getAsyncDedicatedConnection())
+					.from(RedisSortedSetAsyncCommands::bzpopmin, preciseTimeout(timeout, unit), key)
+					.get(it -> it.map(LettuceConverters::toTuple).getValueOrElse(null));
+		}
+
 		return connection.invoke(connection.getAsyncDedicatedConnection())
 				.from(RedisSortedSetAsyncCommands::bzpopmin, unit.toSeconds(timeout), key)
 				.get(it -> it.map(LettuceConverters::toTuple).getValueOrElse(null));
@@ -359,6 +366,13 @@ public Tuple bZPopMax(byte[] key, long timeout, TimeUnit unit) {
 		Assert.notNull(key, "Key must not be null!");
 		Assert.notNull(unit, "TimeUnit must not be null!");
 
+		if(TimeUnit.MILLISECONDS == unit) {
+
+			return connection.invoke(connection.getAsyncDedicatedConnection())
+					.from(RedisSortedSetAsyncCommands::bzpopmax, preciseTimeout(timeout, unit), key)
+					.get(it -> it.map(LettuceConverters::toTuple).getValueOrElse(null));
+		}
+
 		return connection.invoke(connection.getAsyncDedicatedConnection())
 				.from(RedisSortedSetAsyncCommands::bzpopmax, unit.toSeconds(timeout), key)
 				.get(it -> it.map(LettuceConverters::toTuple).getValueOrElse(null));
@@ -703,4 +717,7 @@ private static io.lettuce.core.ZAddArgs toZAddArgs(ZAddArgs source) {
 		return target;
 	}
 
+	static double preciseTimeout(long val, TimeUnit unit) {
+		return (double) unit.toMillis(val) / 1000.0D;
+	}
 }

src/main/java/org/springframework/data/redis/core/ZSetOperations.java

@@ -365,8 +365,8 @@ static <V> TypedTuple<V> of(V value, @Nullable Double score) {
 	Set<TypedTuple<V>> popMin(K key, long count);
 
 	/**
-	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout
@@ -379,8 +379,8 @@ static <V> TypedTuple<V> of(V value, @Nullable Double score) {
 	TypedTuple<V> popMin(K key, long timeout, TimeUnit unit);
 
 	/**
-	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the lowest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout must not be {@literal null}.
@@ -422,8 +422,8 @@ default TypedTuple<V> popMin(K key, Duration timeout) {
 	Set<TypedTuple<V>> popMax(K key, long count);
 
 	/**
-	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout
@@ -436,8 +436,8 @@ default TypedTuple<V> popMin(K key, Duration timeout) {
 	TypedTuple<V> popMax(K key, long timeout, TimeUnit unit);
 
 	/**
-	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <b>Blocks
-	 * connection</b> until element available or {@code timeout} reached.
+	 * Remove and return the value with its score having the highest score from sorted set at {@code key}. <br />
+	 * <b>Blocks connection</b> until element available or {@code timeout} reached.
 	 *
 	 * @param key must not be {@literal null}.
 	 * @param timeout must not be {@literal null}.

src/test/java/org/springframework/data/redis/connection/DefaultStringRedisConnectionTxTests.java

@@ -1318,7 +1318,7 @@ public void testZScore() {
 
 	@Test
 	public void testZMScore() {
-		
+
 		doReturn(Collections.singletonList(Arrays.asList(1d, 3d))).when(nativeConnection).exec();
 		super.testZMScore();
 	}