shantanukawale-stripe commented on code in PR #24028:
URL: https://github.com/apache/kafka/pull/24028#discussion_r4452381176


##########
storage/src/main/java/org/apache/kafka/storage/internals/log/UnifiedLog.java:
##########
@@ -1812,12 +1812,32 @@ public boolean isEmpty() {
         return logStartOffset == logEndOffset();
     }
 
-    private Optional<FileRecords.TimestampAndOffset> 
searchOffsetInLocalLog(long targetTimestamp, long startOffset) throws 
IOException {
+    /**
+     * Find the first local record at or after startOffset whose timestamp is 
at least targetTimestamp.
+     */
+    public Optional<FileRecords.TimestampAndOffset> 
searchOffsetInLocalLog(long targetTimestamp, long startOffset) throws 
IOException {

Review Comment:
   Done, made the two-argument helper private again.



##########
storage/src/main/java/org/apache/kafka/storage/internals/log/UnifiedLog.java:
##########
@@ -1812,12 +1812,32 @@ public boolean isEmpty() {
         return logStartOffset == logEndOffset();
     }
 
-    private Optional<FileRecords.TimestampAndOffset> 
searchOffsetInLocalLog(long targetTimestamp, long startOffset) throws 
IOException {
+    /**
+     * Find the first local record at or after startOffset whose timestamp is 
at least targetTimestamp.
+     */
+    public Optional<FileRecords.TimestampAndOffset> 
searchOffsetInLocalLog(long targetTimestamp, long startOffset) throws 
IOException {
         // Cache to avoid race conditions.
-        List<LogSegment> segments = logSegments();
-        for (LogSegment segment : segments) {
+        return searchOffsetInLocalLog(logSegments(), targetTimestamp, 
startOffset);
+    }
+
+    /**
+     * Find the first record at or after startOffset in the given snapshot of 
local segments whose timestamp is at
+     * least targetTimestamp. The starting offset may be below the first 
segment.
+     */
+    public Optional<FileRecords.TimestampAndOffset> 
searchOffsetInLocalLog(List<LogSegment> segments, long targetTimestamp,

Review Comment:
   Added the requirement that the segments be sorted by base offset in 
ascending order.



##########
storage/src/test/java/org/apache/kafka/server/log/remote/storage/RemoteLogManagerTest.java:
##########
@@ -1672,15 +1673,144 @@ void testFindOffsetByTimestampWithSegmentNotReady() 
throws IOException, RemoteSt
         leaderEpochFileCache.assign(targetLeaderEpoch, startOffset);
         leaderEpochFileCache.assign(12, 500L);
 
-        doTestFindOffsetByTimestamp(ts, startOffset, targetLeaderEpoch, 
validSegmentEpochs, RemoteLogSegmentState.COPY_SEGMENT_STARTED);
+        setUpRemoteSegmentForFindOffsetByTimestamp(ts, startOffset, 
targetLeaderEpoch, validSegmentEpochs, 
RemoteLogSegmentState.COPY_SEGMENT_STARTED);
 
         Optional<FileRecords.TimestampAndOffset> maybeTimestampAndOffset = 
remoteLogManager.findOffsetByTimestamp(tp, ts, startOffset, 
leaderEpochFileCache);
         assertEquals(Optional.empty(), maybeTimestampAndOffset);
     }
 
-    private void doTestFindOffsetByTimestamp(long ts, long startOffset, int 
targetLeaderEpoch,
-                                             TreeMap<Integer, Long> 
validSegmentEpochs,
-                                             RemoteLogSegmentState state) 
throws IOException, RemoteStorageException {
+    @ParameterizedTest
+    @CsvSource({"false, false", "true, false", "true, true"})
+    void 
testFindOffsetByTimestampSearchesNextRemoteSegmentWhenMatchIsBelowLogStartOffset(boolean
 nextEpoch, boolean spansEpochs) throws Exception {
+        // Register the partition without starting background tasks while the 
lookup mocks are configured.
+        remoteLogManager = spy(remoteLogManager);
+        
doNothing().when(remoteLogManager).doHandleLeaderPartition(eq(leaderTopicIdPartition),
 eq(false));
+        long timestamp = time.milliseconds();
+        long baseOffset = 120L;
+        int leaderEpoch = 10;
+        int laterEpoch = nextEpoch ? leaderEpoch + (spansEpochs ? 2 : 1) : 
leaderEpoch;
+        TreeMap<Integer, Long> firstSegmentEpochs = new 
TreeMap<>(Map.of(leaderEpoch, baseOffset));
+        LeaderEpochFileCache cache = new 
LeaderEpochFileCache(leaderTopicIdPartition.topicPartition(), checkpoint, 
scheduler);
+        cache.assign(leaderEpoch, baseOffset);
+        if (spansEpochs) {
+            firstSegmentEpochs.put(leaderEpoch + 1, baseOffset + 2);
+            cache.assign(leaderEpoch + 1, baseOffset + 2);
+        }
+        if (nextEpoch) {
+            cache.assign(laterEpoch, baseOffset + 3);
+        }
+
+        RemoteLogSegmentMetadata firstSegment = 
setUpRemoteSegmentForFindOffsetByTimestamp(timestamp, baseOffset, leaderEpoch,
+                firstSegmentEpochs, 
RemoteLogSegmentState.COPY_SEGMENT_FINISHED);
+        RemoteLogSegmentMetadata laterSegment = new RemoteLogSegmentMetadata(
+                new RemoteLogSegmentId(leaderTopicIdPartition, 
Uuid.randomUuid()), baseOffset + 3, baseOffset + 5,
+                timestamp + 2, brokerId, timestamp, 1024, Optional.empty(), 
RemoteLogSegmentState.COPY_SEGMENT_FINISHED,
+                new TreeMap<>(Map.of(laterEpoch, baseOffset + 3)));
+        
when(remoteLogMetadataManager.listRemoteLogSegments(eq(leaderTopicIdPartition), 
anyInt()))
+                .thenAnswer(invocation -> {
+                    int epoch = invocation.getArgument(1);
+                    return (nextEpoch
+                            ? (epoch < laterEpoch ? List.of(firstSegment) : 
List.of(laterSegment))
+                            : List.of(firstSegment, laterSegment)).iterator();
+                });
+        // The first segment's maximum belongs to a record below the requested 
starting offset.
+        MemoryRecords firstBatch = MemoryRecords.withRecords(baseOffset, 
Compression.NONE, leaderEpoch,
+                new SimpleRecord(timestamp + 2, new byte[0]), new 
SimpleRecord(timestamp - 1, new byte[0]));
+        MemoryRecords secondBatch = MemoryRecords.withRecords(baseOffset + 2, 
Compression.NONE,
+                spansEpochs ? leaderEpoch + 1 : leaderEpoch, new 
SimpleRecord(timestamp - 1, new byte[0]));
+        ByteBuffer firstSegmentRecords = 
ByteBuffer.allocate(firstBatch.sizeInBytes() + secondBatch.sizeInBytes());
+        firstSegmentRecords.put(firstBatch.buffer()).put(secondBatch.buffer());
+        when(remoteStorageManager.fetchLogSegment(firstSegment, 
0)).thenAnswer(invocation ->
+                new ByteArrayInputStream(firstSegmentRecords.array()));
+        when(remoteStorageManager.fetchLogSegment(laterSegment, 
0)).thenAnswer(invocation ->
+                new ByteArrayInputStream(records(timestamp, baseOffset + 3, 
laterEpoch).buffer().array()));
+
+        Optional<FileRecords.TimestampAndOffset> result = 
remoteLogManager.findOffsetByTimestamp(
+                leaderTopicIdPartition.topicPartition(), timestamp, baseOffset 
+ 1, cache);
+        verify(remoteStorageManager).fetchLogSegment(firstSegment, 0);
+        verify(remoteStorageManager).fetchLogSegment(laterSegment, 0);
+        assertEquals(Optional.of(new FileRecords.TimestampAndOffset(timestamp 
+ 1, baseOffset + 4, Optional.of(laterEpoch))), result);
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void 
testFindOffsetByTimestampSearchesNextLocalSegmentWhenMatchIsBelowLogStartOffset(boolean
 laterMatch) throws Exception {
+        long timestamp = time.milliseconds();
+        int leaderEpoch = 10;
+        LeaderEpochFileCache cache = new 
LeaderEpochFileCache(leaderTopicIdPartition.topicPartition(), checkpoint, 
scheduler);
+        cache.assign(leaderEpoch, 0L);
+        RemoteLogSegmentMetadata metadata = new RemoteLogSegmentMetadata(
+                new RemoteLogSegmentId(leaderTopicIdPartition, 
Uuid.randomUuid()), 0L, 1L,
+                timestamp + 1, brokerId, timestamp, 1024, Optional.empty(), 
RemoteLogSegmentState.COPY_SEGMENT_FINISHED,
+                new TreeMap<>(Map.of(leaderEpoch, 0L)));
+        RemoteLogSegmentMetadata laterMetadata = new RemoteLogSegmentMetadata(
+                new RemoteLogSegmentId(leaderTopicIdPartition, 
Uuid.randomUuid()), 2L, 2L,
+                timestamp, brokerId, timestamp, 1024, Optional.empty(), 
RemoteLogSegmentState.COPY_SEGMENT_FINISHED,
+                new TreeMap<>(Map.of(leaderEpoch, 2L)));
+        
when(remoteLogMetadataManager.listRemoteLogSegments(leaderTopicIdPartition, 
leaderEpoch))
+                .thenAnswer(invocation -> List.of(metadata, 
laterMetadata).iterator());
+        when(mockLog.logEndOffset()).thenReturn(3L);
+        remoteLogManager.onLeadershipChange(Set.of(), 
Set.of(mockPartition(leaderTopicIdPartition)), topicIds);
+
+        try (LogSegment firstSegment = spy(LogSegment.open(new File(logDir), 
0L, mockLog.config(), time, 0, false));
+             LogSegment laterSegment = LogSegment.open(new File(logDir), 2L, 
mockLog.config(), time, 0, false)) {
+            firstSegment.append(1L, MemoryRecords.withRecords(0L, 
Compression.NONE, leaderEpoch,
+                    new SimpleRecord(timestamp + 1, new byte[0]), new 
SimpleRecord(timestamp - 1, new byte[0])));
+            laterSegment.append(2L, MemoryRecords.withRecords(2L, 
Compression.NONE, leaderEpoch,
+                    new SimpleRecord(laterMatch ? timestamp : timestamp - 1, 
new byte[0])));
+            when(mockLog.logSegments()).thenReturn(List.of(firstSegment, 
laterSegment));
+            when(mockLog.searchOffsetInLocalLog(anyList(), anyLong(), 
anyLong())).thenCallRealMethod();
+
+            Optional<FileRecords.TimestampAndOffset> expected = laterMatch
+                    ? Optional.of(new 
FileRecords.TimestampAndOffset(timestamp, 2L, Optional.of(leaderEpoch)))
+                    : Optional.empty();
+            assertEquals(expected, remoteLogManager.findOffsetByTimestamp(
+                    leaderTopicIdPartition.topicPartition(), timestamp, 1L, 
cache));
+            verify(firstSegment).findOffsetByTimestamp(timestamp, 1L, 
mockLog.config().maxDecompressedMessageBytes());
+            verify(mockLog).searchOffsetInLocalLog(List.of(firstSegment, 
laterSegment), timestamp, 1L);
+            verify(remoteStorageManager, never()).fetchLogSegment(any(), 
anyInt());
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void testFindOffsetByTimestampUsesLocalSegmentSnapshot(boolean laterMatch) 
throws Exception {
+        remoteLogManager = spy(remoteLogManager);
+        
doNothing().when(remoteLogManager).doHandleFollowerPartition(leaderTopicIdPartition);

Review Comment:
   Added a comment explaining why: the test needs to register the partition 
without starting the follower's background remote-offset updates against its 
mocks.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to