Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c656ecf

Browse files
noojungnoojung
committed
Enhance RedisItemReader/Writer performance with pipelining
Signed-off-by: Hyunwoo Jung <[email protected]>
1 parent 3bcc525 commit c656ecf

File tree

9 files changed

+193
-51
lines changed

9 files changed

+193
-51
lines changed

‎spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemReader.java‎

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,21 +18,27 @@
1818
import org.springframework.batch.item.ExecutionContext;
1919
import org.springframework.batch.item.ItemStreamException;
2020
import org.springframework.batch.item.ItemStreamReader;
21-
import org.springframework.data.redis.core.Cursor;
22-
import org.springframework.data.redis.core.RedisTemplate;
23-
import org.springframework.data.redis.core.ScanOptions;
21+
import org.springframework.dao.DataAccessException;
22+
import org.springframework.data.redis.core.*;
2423
import org.springframework.util.Assert;
2524

25+
import java.util.ArrayDeque;
26+
import java.util.ArrayList;
27+
import java.util.Deque;
28+
import java.util.List;
29+
2630
/**
2731
* Item reader for Redis based on Spring Data Redis. Uses a {@link RedisTemplate} to query
2832
* data. The user should provide a {@link ScanOptions} to specify the set of keys to
29-
* query.
33+
* query. The {@code fetchSize} property controls how many items are fetched from Redis in
34+
* a single pipeline round-trip for efficiency.
3035
*
3136
* <p>
3237
* The implementation is not thread-safe and not restartable.
3338
* </p>
3439
*
3540
* @author Mahmoud Ben Hassine
41+
* @author Hyunwoo Jung
3642
* @since 5.1
3743
* @param <K> type of keys
3844
* @param <V> type of values
@@ -43,13 +49,20 @@ public class RedisItemReader<K, V> implements ItemStreamReader<V> {
4349

4450
private final ScanOptions scanOptions;
4551

52+
private final int fetchSize;
53+
54+
private final Deque<V> buffer;
55+
4656
private Cursor<K> cursor;
4757

48-
public RedisItemReader(RedisTemplate<K, V> redisTemplate, ScanOptions scanOptions) {
58+
public RedisItemReader(RedisTemplate<K, V> redisTemplate, ScanOptions scanOptions, int fetchSize) {
4959
Assert.notNull(redisTemplate, "redisTemplate must not be null");
5060
Assert.notNull(scanOptions, "scanOptions must no be null");
61+
Assert.isTrue(fetchSize > 0, "fetchSize must be greater than 0");
5162
this.redisTemplate = redisTemplate;
5263
this.scanOptions = scanOptions;
64+
this.fetchSize = fetchSize;
65+
this.buffer = new ArrayDeque<>();
5366
}
5467

5568
@Override
@@ -59,18 +72,45 @@ public void open(ExecutionContext executionContext) throws ItemStreamException {
5972

6073
@Override
6174
public V read() throws Exception {
62-
if (this.cursor.hasNext()) {
63-
K nextKey = this.cursor.next();
64-
return this.redisTemplate.opsForValue().get(nextKey);
65-
}
66-
else {
67-
return null;
75+
if (this.buffer.isEmpty()) {
76+
fetchNext();
6877
}
78+
79+
return this.buffer.pollFirst();
6980
}
7081

7182
@Override
7283
public void close() throws ItemStreamException {
7384
this.cursor.close();
7485
}
7586

87+
private void fetchNext() {
88+
List<K> keys = new ArrayList<>();
89+
while (this.cursor.hasNext() && keys.size() < this.fetchSize) {
90+
keys.add(this.cursor.next());
91+
}
92+
93+
if (keys.isEmpty()) {
94+
return;
95+
}
96+
97+
@SuppressWarnings("unchecked")
98+
List<V> items = (List<V>) this.redisTemplate.executePipelined(sessionCallback(keys));
99+
100+
this.buffer.addAll(items);
101+
}
102+
103+
private SessionCallback<Object> sessionCallback(List<K> keys) {
104+
return new SessionCallback<>() {
105+
@Override
106+
public Object execute(RedisOperations operations) throws DataAccessException {
107+
for (K key : keys) {
108+
operations.opsForValue().get(key);
109+
}
110+
111+
return null;
112+
}
113+
};
114+
}
115+
76116
}

‎spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemWriter.java‎

Lines changed: 52 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,37 +18,53 @@
1818

1919
import org.springframework.batch.item.ItemWriter;
2020
import org.springframework.batch.item.KeyValueItemWriter;
21+
import org.springframework.dao.DataAccessException;
22+
import org.springframework.data.redis.core.RedisOperations;
2123
import org.springframework.data.redis.core.RedisTemplate;
24+
import org.springframework.data.redis.core.SessionCallback;
25+
import org.springframework.data.util.Pair;
2226
import org.springframework.util.Assert;
2327

28+
import java.util.ArrayList;
29+
import java.util.List;
30+
2431
/**
2532
* <p>
2633
* An {@link ItemWriter} implementation for Redis using a {@link RedisTemplate} .
2734
* </p>
2835
*
2936
* @author Santiago Molano
3037
* @author Mahmoud Ben Hassine
38+
* @author Hyunwoo Jung
3139
* @since 5.1
3240
*/
3341
public class RedisItemWriter<K, T> extends KeyValueItemWriter<K, T> {
3442

3543
private RedisTemplate<K, T> redisTemplate;
3644

45+
private final List<Pair<K, T>> buffer = new ArrayList<>();
46+
3747
@Override
3848
protected void writeKeyValue(K key, T value) {
39-
if (this.delete) {
40-
this.redisTemplate.delete(key);
41-
}
42-
else {
43-
this.redisTemplate.opsForValue().set(key, value);
44-
}
49+
this.buffer.add(Pair.of(key, value));
4550
}
4651

4752
@Override
4853
protected void init() {
4954
Assert.notNull(this.redisTemplate, "RedisTemplate must not be null");
5055
}
5156

57+
@Override
58+
protected void flush() throws Exception {
59+
if (this.buffer.isEmpty()) {
60+
return;
61+
}
62+
63+
this.redisTemplate.executePipelined(sessionCallback());
64+
65+
this.buffer.clear();
66+
}
67+
5268
/**
5369
* Set the {@link RedisTemplate} to use.
5470
* @param redisTemplate the template to use
@@ -57,4 +73,33 @@ public void setRedisTemplate(RedisTemplate<K, T> redisTemplate) {
5773
this.redisTemplate = redisTemplate;
5874
}
5975

76+
private SessionCallback<Object> sessionCallback() {
77+
return new SessionCallback<>() {
78+
79+
@SuppressWarnings("unchecked")
80+
@Override
81+
public Object execute(RedisOperations operations) throws DataAccessException {
82+
if (RedisItemWriter.this.delete) {
83+
executeDeleteOperations(operations);
84+
}
85+
else {
86+
executeSetOperations(operations);
87+
}
88+
return null;
89+
}
90+
};
91+
}
92+
93+
private void executeDeleteOperations(RedisOperations<K, T> operations) {
94+
for (Pair<K, T> item : this.buffer) {
95+
operations.delete(item.getFirst());
96+
}
97+
}
98+
99+
private void executeSetOperations(RedisOperations<K, T> operations) {
100+
for (Pair<K, T> item : this.buffer) {
101+
operations.opsForValue().set(item.getFirst(), item.getSecond());
102+
}
103+
}
104+
60105
}

‎spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilder.java‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
2323
* Builder for {@link RedisItemReader}.
2424
*
2525
* @author Mahmoud Ben Hassine
26+
* @author Hyunwoo Jung
2627
* @since 5.1
2728
* @param <K> type of keys
2829
* @param <V> type of values
@@ -33,6 +34,8 @@ public class RedisItemReaderBuilder<K, V> {
3334

3435
private ScanOptions scanOptions;
3536

37+
private int fetchSize;
38+
3639
/**
3740
* Set the {@link RedisTemplate} to use in the reader.
3841
* @param redisTemplate the template to use
@@ -53,12 +56,22 @@ public RedisItemReaderBuilder<K, V> scanOptions(ScanOptions scanOptions) {
5356
return this;
5457
}
5558

59+
/**
60+
* Set the fetchSize to how many items from Redis in a single pipeline round-trip.
61+
* @param fetchSize the number of items to fetch per pipeline execution
62+
* @return the current builder instance for fluent chaining
63+
*/
64+
public RedisItemReaderBuilder<K, V> fetchSize(int fetchSize) {
65+
this.fetchSize = fetchSize;
66+
return this;
67+
}
68+
5669
/**
5770
* Build a new {@link RedisItemReader}.
5871
* @return a new item reader
5972
*/
6073
public RedisItemReader<K, V> build() {
61-
return new RedisItemReader<>(this.redisTemplate, this.scanOptions);
74+
return new RedisItemReader<>(this.redisTemplate, this.scanOptions, this.fetchSize);
6275
}
6376

6477
}

‎spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderIntegrationTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ void testRead(RedisConnectionFactory connectionFactory) throws Exception {
8282

8383
RedisTemplate<String, Person> redisTemplate = setUpRedisTemplate(connectionFactory);
8484
ScanOptions scanOptions = ScanOptions.scanOptions().match("person:*").count(10).build();
85-
this.reader = new RedisItemReader<>(redisTemplate, scanOptions);
85+
this.reader = new RedisItemReader<>(redisTemplate, scanOptions, 10);
8686

8787
this.reader.open(new ExecutionContext());
8888

‎spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderTests.java‎

Lines changed: 41 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -16,20 +16,30 @@
1616
package org.springframework.batch.item.redis;
1717

1818
import org.junit.jupiter.api.Assertions;
19+
import org.junit.jupiter.api.BeforeEach;
1920
import org.junit.jupiter.api.Test;
2021
import org.junit.jupiter.api.extension.ExtendWith;
2122
import org.mockito.Answers;
2223
import org.mockito.Mock;
23-
import org.mockito.Mockito;
2424
import org.mockito.junit.jupiter.MockitoExtension;
25-
2625
import org.springframework.batch.item.ExecutionContext;
2726
import org.springframework.data.redis.core.Cursor;
2827
import org.springframework.data.redis.core.RedisTemplate;
2928
import org.springframework.data.redis.core.ScanOptions;
29+
import org.springframework.data.redis.core.SessionCallback;
30+
31+
import java.util.ArrayList;
32+
import java.util.List;
33+
34+
import static org.mockito.ArgumentMatchers.any;
35+
import static org.mockito.Mockito.when;
3036

37+
/**
38+
* @author Mahmoud Ben Hassine
39+
* @author Hyunwoo Jung
40+
*/
3141
@ExtendWith(MockitoExtension.class)
32-
public class RedisItemReaderTests {
42+
class RedisItemReaderTests {
3343

3444
@Mock(answer = Answers.RETURNS_DEEP_STUBS)
3545
private RedisTemplate<String, String> redisTemplate;
@@ -40,15 +50,36 @@ public class RedisItemReaderTests {
4050
@Mock
4151
private Cursor<String> cursor;
4252

53+
private List<String> results;
54+
55+
@BeforeEach
56+
void setUp() {
57+
this.results = new ArrayList<>();
58+
59+
when(this.redisTemplate.executePipelined(any(SessionCallback.class))).thenAnswer(invocation -> {
60+
SessionCallback<?> sessionCallback = invocation.getArgument(0);
61+
sessionCallback.execute(this.redisTemplate);
62+
return this.results;
63+
});
64+
}
65+
4366
@Test
4467
void testRead() throws Exception {
4568
// given
46-
Mockito.when(this.redisTemplate.scan(this.scanOptions)).thenReturn(this.cursor);
47-
Mockito.when(this.cursor.hasNext()).thenReturn(true, true, false);
48-
Mockito.when(this.cursor.next()).thenReturn("person:1", "person:2");
49-
Mockito.when(this.redisTemplate.opsForValue().get("person:1")).thenReturn("foo");
50-
Mockito.when(this.redisTemplate.opsForValue().get("person:2")).thenReturn("bar");
51-
RedisItemReader<String, String> redisItemReader = new RedisItemReader<>(this.redisTemplate, this.scanOptions);
69+
when(this.redisTemplate.scan(this.scanOptions)).thenReturn(this.cursor);
70+
when(this.cursor.hasNext()).thenReturn(true, true, false);
71+
when(this.cursor.next()).thenReturn("person:1", "person:2");
72+
when(this.redisTemplate.opsForValue().get("person:1")).thenAnswer(invocation -> {
73+
results.add("foo");
74+
return null;
75+
});
76+
when(this.redisTemplate.opsForValue().get("person:2")).thenAnswer(invocation -> {
77+
results.add("bar");
78+
return null;
79+
});
80+
81+
RedisItemReader<String, String> redisItemReader = new RedisItemReader<>(this.redisTemplate, this.scanOptions,
82+
10);
5283
redisItemReader.open(new ExecutionContext());
5384

5485
// when

‎spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemWriterIntegrationTests.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@
3939

4040
import java.util.stream.Stream;
4141

42-
import static org.junit.jupiter.api.Assertions.*;
42+
import static org.junit.jupiter.api.Assertions.assertEquals;
43+
import static org.junit.jupiter.api.Assertions.assertFalse;
4344

4445
/**
4546
* @author Hyunwoo Jung

0 commit comments

Comments
 (0)