Skip to content

Commit 6d0e836

Browse files
bdshadowbdshadow
committed
perf: pipeline Redis batch-job state cleanup with RBatch
removeJobState now deletes the state hash, the 7 counters and both markers in a single RBatch instead of ~10 sequential round-trips. clearUnusedStates reads every job's state in one pipelined batch (chunked to bound memory and pipeline size) and deletes the completed ones in a second batch, replacing the previous per-job round-trips. When a pipelined decode fails it falls back to isolated per-job reads so one unreadable job cannot block cleanup of the rest.
1 parent 0f5d47d commit 6d0e836

2 files changed

Lines changed: 237 additions & 34 deletions

File tree

‎backend/app/src/test/kotlin/io/tolgee/batch/state/RedisBatchJobStateStorageTest.kt‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
package io.tolgee.batch.state
2+
3+
import io.tolgee.AbstractSpringTest
4+
import io.tolgee.fixtures.RedisRunner
5+
import io.tolgee.model.batch.BatchJobChunkExecutionStatus
6+
import io.tolgee.testing.ContextRecreatingTest
7+
import io.tolgee.testing.assert
8+
import org.junit.jupiter.api.AfterAll
9+
import org.junit.jupiter.api.AfterEach
10+
import org.junit.jupiter.api.Test
11+
import org.mockito.kotlin.any
12+
import org.mockito.kotlin.clearInvocations
13+
import org.mockito.kotlin.never
14+
import org.mockito.kotlin.times
15+
import org.mockito.kotlin.verify
16+
import org.redisson.api.RedissonClient
17+
import org.redisson.client.codec.StringCodec
18+
import org.springframework.beans.factory.annotation.Autowired
19+
import org.springframework.boot.test.context.SpringBootTest
20+
import org.springframework.boot.test.util.TestPropertyValues
21+
import org.springframework.context.ApplicationContextInitializer
22+
import org.springframework.context.ConfigurableApplicationContext
23+
import org.springframework.test.annotation.DirtiesContext
24+
import org.springframework.test.context.ContextConfiguration
25+
import org.springframework.test.context.bean.override.mockito.MockitoSpyBean
26+
27+
@ContextRecreatingTest
28+
@SpringBootTest(
29+
properties = [
30+
"tolgee.cache.use-redis=true",
31+
"tolgee.cache.enabled=true",
32+
"tolgee.websocket.use-redis=true",
33+
],
34+
)
35+
@ContextConfiguration(initializers = [RedisBatchJobStateStorageTest.Companion.Initializer::class])
36+
@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_CLASS)
37+
class RedisBatchJobStateStorageTest : AbstractSpringTest() {
38+
companion object {
39+
private const val STATE_KEY_PREFIX = "batch_job_state:"
40+
private const val STATE_INITIALIZED_PREFIX = "batch_job_state_initialized:"
41+
private const val STARTED_PREFIX = "batch_job_started:"
42+
43+
val redisRunner = RedisRunner()
44+
45+
@AfterAll
46+
@JvmStatic
47+
fun stopRedis() {
48+
redisRunner.stop()
49+
}
50+
51+
class Initializer : ApplicationContextInitializer<ConfigurableApplicationContext> {
52+
override fun initialize(configurableApplicationContext: ConfigurableApplicationContext) {
53+
redisRunner.run()
54+
TestPropertyValues
55+
.of("spring.data.redis.port=${RedisRunner.port}")
56+
.applyTo(configurableApplicationContext)
57+
}
58+
}
59+
}
60+
61+
@Autowired
62+
lateinit var stateProvider: BatchJobStateProvider
63+
64+
@MockitoSpyBean
65+
@Autowired
66+
lateinit var redissonClient: RedissonClient
67+
68+
private val usedJobIds = mutableSetOf<Long>()
69+
70+
@AfterEach
71+
fun cleanup() {
72+
usedJobIds.forEach { stateProvider.removeJobState(it) }
73+
usedJobIds.clear()
74+
}
75+
76+
@Test
77+
fun `removeJobState pipelines its deletes into a single RBatch`() {
78+
val jobId =
79+
seed(900_020L, BatchJobChunkExecutionStatus.SUCCESS) {
80+
incrementRunningCount(it)
81+
tryMarkJobStarted(it)
82+
}
83+
84+
clearInvocations(redissonClient)
85+
stateProvider.removeJobState(jobId)
86+
87+
// The deletes must be pipelined through one RBatch, not issued as per-key round-trips.
88+
verify(redissonClient, times(1)).createBatch()
89+
verify(redissonClient, never()).getAtomicLong(any<String>())
90+
91+
stateProvider.hasCachedJobState(jobId).assert.isFalse()
92+
}
93+
94+
@Test
95+
fun `clearUnusedStates clears completed jobs across chunks and keeps counters and unfinished jobs`() {
96+
// Exceeds the production CLEANUP_BATCH_SIZE (100) so cleanup runs across more than one chunk.
97+
val completedJobIds = (901_000L until 901_150L).toList()
98+
completedJobIds.forEach {
99+
seed(it, BatchJobChunkExecutionStatus.SUCCESS) {
100+
incrementRunningCount(it)
101+
tryMarkJobStarted(it)
102+
}
103+
}
104+
val runningJobId = seed(900_002L, BatchJobChunkExecutionStatus.RUNNING)
105+
106+
stateProvider.clearUnusedStates()
107+
108+
// Every completed job's hash is cleared — including those in the second chunk.
109+
completedJobIds.forEach { stateProvider.hasCachedJobState(it).assert.isFalse() }
110+
111+
// Counters and the started marker survive until removeJobState finalizes the job
112+
// (dropping the started marker would allow re-execution).
113+
val sample = completedJobIds.first()
114+
isInitializedBucketPresent(sample).assert.isFalse()
115+
stateProvider.getRunningCount(sample).assert.isEqualTo(1)
116+
isStartedBucketPresent(sample).assert.isTrue()
117+
118+
// Unfinished job is left untouched.
119+
stateProvider.hasCachedJobState(runningJobId).assert.isTrue()
120+
isInitializedBucketPresent(runningJobId).assert.isTrue()
121+
}
122+
123+
@Test
124+
fun `clearUnusedStates skips a job whose state cannot be read and still cleans healthy jobs`() {
125+
val healthyJobId = seed(900_040L, BatchJobChunkExecutionStatus.SUCCESS)
126+
val corruptJobId = 900_041L
127+
usedJobIds += corruptJobId
128+
// Forge a value the state codec cannot deserialize, so this job's read future fails.
129+
redissonClient.getMap<String, String>("$STATE_KEY_PREFIX$corruptJobId", StringCodec.INSTANCE)["1"] = "corrupt"
130+
redissonClient.getBucket<Boolean>("$STATE_INITIALIZED_PREFIX$corruptJobId").set(true)
131+
132+
stateProvider.clearUnusedStates()
133+
134+
// The unreadable job is skipped, not deleted...
135+
stateProvider.hasCachedJobState(corruptJobId).assert.isTrue()
136+
// ...while the healthy completed job in the same pass is still cleaned.
137+
stateProvider.hasCachedJobState(healthyJobId).assert.isFalse()
138+
}
139+
140+
private fun seed(
141+
jobId: Long,
142+
status: BatchJobChunkExecutionStatus,
143+
counters: BatchJobStateProvider.(Long) -> Unit = {},
144+
): Long {
145+
usedJobIds += jobId
146+
stateProvider.updateSingleExecution(jobId, 1L, executionState(status))
147+
redissonClient.getBucket<Boolean>("$STATE_INITIALIZED_PREFIX$jobId").set(true)
148+
stateProvider.counters(jobId)
149+
return jobId
150+
}
151+
152+
private fun executionState(status: BatchJobChunkExecutionStatus) =
153+
ExecutionState(
154+
successTargetsCount = 0,
155+
status = status,
156+
chunkNumber = 0,
157+
retry = null,
158+
transactionCommitted = true,
159+
)
160+
161+
private fun isInitializedBucketPresent(jobId: Long) =
162+
redissonClient.getBucket<Boolean>("$STATE_INITIALIZED_PREFIX$jobId").isExists
163+
164+
private fun isStartedBucketPresent(jobId: Long) = redissonClient.getBucket<Boolean>("$STARTED_PREFIX$jobId").isExists
165+
}

‎backend/data/src/main/kotlin/io/tolgee/batch/state/RedisBatchJobStateStorage.kt‎

Lines changed: 72 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import io.tolgee.model.batch.BatchJobChunkExecution
55
import io.tolgee.util.Logging
66
import io.tolgee.util.logger
77
import org.redisson.api.RAtomicLong
8+
import org.redisson.api.RBatch
89
import org.redisson.api.RMap
910
import org.redisson.api.RedissonClient
1011
import java.util.concurrent.ConcurrentHashMap
@@ -45,6 +46,8 @@ open class RedisBatchJobStateStorage(
4546
private const val REDIS_CANCELLED_COUNT_KEY_PREFIX = "batch_job_cancelled:"
4647
private const val REDIS_COMMITTED_COUNT_KEY_PREFIX = "batch_job_committed:"
4748
private const val REDIS_STARTED_KEY_PREFIX = "batch_job_started:"
49+
50+
private const val CLEANUP_BATCH_SIZE = 100
4851
}
4952

5053
// Local cache for initialization status - avoids Redis calls for already-initialized jobs
@@ -175,7 +178,7 @@ open class RedisBatchJobStateStorage(
175178
}
176179

177180
override fun tryMarkJobStarted(jobId: Long): Boolean {
178-
val bucket = redissonClient.getBucket<Boolean>("$REDIS_STARTED_KEY_PREFIX$jobId")
181+
val bucket = redissonClient.getBucket<Boolean>(startedKey(jobId))
179182
return bucket.setIfAbsent(true)
180183
}
181184

@@ -192,13 +195,12 @@ open class RedisBatchJobStateStorage(
192195

193196
override fun removeJobState(jobId: Long) {
194197
logger.debug("Removing job state for job $jobId")
195-
removeAllCounters(jobId)
196-
val redisHash = getRedisHashForJob(jobId)
197-
redisHash.delete()
198-
// Also remove initialization and started markers
199-
redissonClient.getBucket<Boolean>("$REDIS_STATE_INITIALIZED_KEY_PREFIX$jobId").delete()
200-
redissonClient.getBucket<Boolean>("$REDIS_STARTED_KEY_PREFIX$jobId").delete()
201-
// Clear local initialization cache to allow re-initialization if jobId is reused
198+
val batch = redissonClient.createBatch()
199+
addCounterDeletesToBatch(batch, jobId)
200+
batch.getMap<Long, ExecutionState>(stateKey(jobId)).deleteAsync()
201+
batch.getBucket<Boolean>(initializedKey(jobId)).deleteAsync()
202+
batch.getBucket<Boolean>(startedKey(jobId)).deleteAsync()
203+
batch.execute()
202204
localInitializedJobs.remove(jobId)
203205
}
204206

@@ -213,25 +215,52 @@ open class RedisBatchJobStateStorage(
213215

214216
/**
215217
* Cleans up batch job state hashes where all executions have a completed status.
216-
*
217-
* This uses HVALS to read all values from each hash. Because [ExecutionState] only
218-
* stores lightweight metadata (no successTargets list), each value is ~50 bytes,
219-
* making this operation fast even for jobs with thousands of chunks.
220218
*/
221219
override fun clearUnusedStates() {
222-
val keys = redissonClient.keys.getKeysByPattern("$REDIS_STATE_KEY_PREFIX*")
223-
keys.forEach { key ->
224-
val jobId = key.removePrefix(REDIS_STATE_KEY_PREFIX).toLongOrNull() ?: return@forEach
225-
val redisHash = getRedisHashForJob(jobId)
226-
val allCompleted = redisHash.readAllValues().all { state -> state.status.completed }
227-
if (allCompleted) {
228-
redisHash.delete()
229-
redissonClient.getBucket<Boolean>("$REDIS_STATE_INITIALIZED_KEY_PREFIX$jobId").delete()
230-
// Clear local initialization cache to allow re-initialization if jobId is reused
231-
localInitializedJobs.remove(jobId)
232-
// Do NOT remove counters here - they're needed until job status is properly updated
233-
// Counters will be removed in removeJobState when job is finalized
220+
getCachedJobIds().chunked(CLEANUP_BATCH_SIZE).forEach { clearCompletedStates(it) }
221+
}
222+
223+
private fun clearCompletedStates(jobIds: List<Long>) {
224+
val completedJobIds =
225+
readStates(jobIds).mapNotNull { (jobId, values) ->
226+
val allCompleted = values != null && values.all { it.status.completed }
227+
jobId.takeIf { allCompleted }
234228
}
229+
if (completedJobIds.isEmpty()) {
230+
return
231+
}
232+
233+
// Counters stay live until removeJobState finalizes the job; deleting them here corrupts status updates.
234+
val deleteBatch = redissonClient.createBatch()
235+
completedJobIds.forEach { jobId ->
236+
deleteBatch.getMap<Long, ExecutionState>(stateKey(jobId)).deleteAsync()
237+
deleteBatch.getBucket<Boolean>(initializedKey(jobId)).deleteAsync()
238+
}
239+
deleteBatch.execute()
240+
completedJobIds.forEach { localInitializedJobs.remove(it) }
241+
}
242+
243+
private fun readStates(jobIds: List<Long>): Map<Long, Collection<ExecutionState>?> {
244+
return try {
245+
val readBatch = redissonClient.createBatch()
246+
val futures =
247+
jobIds.associateWith { jobId ->
248+
readBatch.getMap<Long, ExecutionState>(stateKey(jobId)).readAllValuesAsync()
249+
}
250+
readBatch.execute()
251+
futures.mapValues { it.value.get() }
252+
} catch (e: Exception) {
253+
logger.warn("Pipelined batch-job state read failed; falling back to per-job reads", e)
254+
jobIds.associateWith { readStateOrNull(it) }
255+
}
256+
}
257+
258+
private fun readStateOrNull(jobId: Long): Collection<ExecutionState>? {
259+
return try {
260+
getRedisHashForJob(jobId).readAllValues()
261+
} catch (e: Exception) {
262+
logger.warn("Failed to read batch job state for job $jobId during cleanup", e)
263+
null
235264
}
236265
}
237266

@@ -251,9 +280,15 @@ open class RedisBatchJobStateStorage(
251280
}
252281

253282
private fun getRedisHashForJob(jobId: Long): RMap<Long, ExecutionState> {
254-
return redissonClient.getMap("$REDIS_STATE_KEY_PREFIX$jobId")
283+
return redissonClient.getMap(stateKey(jobId))
255284
}
256285

286+
private fun stateKey(jobId: Long) = "$REDIS_STATE_KEY_PREFIX$jobId"
287+
288+
private fun initializedKey(jobId: Long) = "$REDIS_STATE_INITIALIZED_KEY_PREFIX$jobId"
289+
290+
private fun startedKey(jobId: Long) = "$REDIS_STARTED_KEY_PREFIX$jobId"
291+
257292
/**
258293
* Ensures Redis hash is initialized from DB. Uses local in-memory cache first for O(1) check,
259294
* then falls back to Redis marker for cross-instance coordination.
@@ -270,7 +305,7 @@ open class RedisBatchJobStateStorage(
270305
return
271306
}
272307
// Check Redis marker for cross-instance coordination
273-
val initKey = "$REDIS_STATE_INITIALIZED_KEY_PREFIX$jobId"
308+
val initKey = initializedKey(jobId)
274309
if (redissonClient.getBucket<Boolean>(initKey).get() == true) {
275310
localInitializedJobs.add(jobId)
276311
return
@@ -295,14 +330,17 @@ open class RedisBatchJobStateStorage(
295330
}
296331
}
297332

298-
private fun removeAllCounters(jobId: Long) {
299-
redissonClient.getAtomicLong("$REDIS_RUNNING_COUNT_KEY_PREFIX$jobId").delete()
300-
redissonClient.getAtomicLong("$REDIS_COMPLETED_CHUNKS_COUNT_KEY_PREFIX$jobId").delete()
301-
redissonClient.getAtomicLong("$REDIS_PROGRESS_COUNT_KEY_PREFIX$jobId").delete()
302-
redissonClient.getAtomicLong("$REDIS_SINGLE_CHUNK_PROGRESS_COUNT_KEY_PREFIX$jobId").delete()
303-
redissonClient.getAtomicLong("$REDIS_FAILED_COUNT_KEY_PREFIX$jobId").delete()
304-
redissonClient.getAtomicLong("$REDIS_CANCELLED_COUNT_KEY_PREFIX$jobId").delete()
305-
redissonClient.getAtomicLong("$REDIS_COMMITTED_COUNT_KEY_PREFIX$jobId").delete()
333+
private fun addCounterDeletesToBatch(
334+
batch: RBatch,
335+
jobId: Long,
336+
) {
337+
batch.getAtomicLong("$REDIS_RUNNING_COUNT_KEY_PREFIX$jobId").deleteAsync()
338+
batch.getAtomicLong("$REDIS_COMPLETED_CHUNKS_COUNT_KEY_PREFIX$jobId").deleteAsync()
339+
batch.getAtomicLong("$REDIS_PROGRESS_COUNT_KEY_PREFIX$jobId").deleteAsync()
340+
batch.getAtomicLong("$REDIS_SINGLE_CHUNK_PROGRESS_COUNT_KEY_PREFIX$jobId").deleteAsync()
341+
batch.getAtomicLong("$REDIS_FAILED_COUNT_KEY_PREFIX$jobId").deleteAsync()
342+
batch.getAtomicLong("$REDIS_CANCELLED_COUNT_KEY_PREFIX$jobId").deleteAsync()
343+
batch.getAtomicLong("$REDIS_COMMITTED_COUNT_KEY_PREFIX$jobId").deleteAsync()
306344
}
307345

308346
private fun initializeCountersFromState(

0 commit comments

Comments
 (0)