From c4d03e04388e7e2b70092835a34f97facabb925c Mon Sep 17 00:00:00 2001
From: Mahmoud Ben Hassine <mbenhassine@vmware.com>
Date: Thu, 10 Oct 2024 10:34:51 +0200
Subject: [PATCH] Pass the current chunk as context in
 ChunkListener#beforeChunk

---
 .../batch/experimental/core/step/item/ChunkListener.java    | 4 ++--
 .../experimental/core/step/item/ChunkOrientedStep.java      | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkListener.java b/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkListener.java
index b019fac..5376ccd 100644
--- a/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkListener.java
+++ b/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkListener.java
@@ -10,12 +10,12 @@
  *
  * Notable difference: afterChunk is called inside the transaction, not outside the transaction.
  */
-public interface ChunkListener<O> extends StepListener {
+public interface ChunkListener<I, O> extends StepListener {
 
 	/**
 	 * Callback before the chunk is processed, inside the transaction.
 	 */
-	default void beforeChunk() {
+	default void beforeChunk(Chunk<I> chunk) {
 	}
 
 	/**
diff --git a/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkOrientedStep.java b/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkOrientedStep.java
index 3dd5e34..d755ff4 100644
--- a/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkOrientedStep.java
+++ b/src/main/java/org/springframework/batch/experimental/core/step/item/ChunkOrientedStep.java
@@ -84,7 +84,7 @@ public boolean rollbackOn(Throwable ex) {
 	 */
 	private final int chunkSize;
 	private final ChunkTracker chunkTracker = new ChunkTracker();
-	private ChunkListener<O> chunkListener = new ChunkListener<>() {};
+	private ChunkListener<I, O> chunkListener = new ChunkListener<>() {};
 
 	/*
 	 * Step state / interruption parameters
@@ -132,8 +132,8 @@ protected void doInTransactionWithoutResult(TransactionStatus status) {
 					Chunk<I> inputChunk = new Chunk<>();
 					Chunk<O> processedChunk = new Chunk<>();
 					try {
-						chunkListener.beforeChunk();
 						inputChunk = read(contribution);
+						chunkListener.beforeChunk(inputChunk);
 						processedChunk = process(inputChunk, contribution);
 						write(processedChunk, contribution);
 						chunkListener.afterChunk(processedChunk);
@@ -243,7 +243,7 @@ public void setItemWriteListener(ItemWriteListener<O> itemWriteListener) {
 		this.itemWriteListener = itemWriteListener;
 	}
 
-	public void setChunkListener(ChunkListener<O> chunkListener) {
+	public void setChunkListener(ChunkListener<I, O> chunkListener) {
 		this.chunkListener = chunkListener;
 	}