Skip to content

Commit

Permalink
Integrate RLMQuotaManager for throttling copies to remote storage
Browse files Browse the repository at this point in the history
  • Loading branch information
@abhijeetk88
abhijeetk88 committed May 24, 2024
1 parent 56549ea commit 3593def
Show file tree
Hide file tree
Showing 4 changed files with 267 additions and 0 deletions.
1 change: 1 addition & 0 deletions checkstyle/import-control-core.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
<allow pkg="kafka.serializer" />
<allow pkg="org.apache.kafka.common" />
<allow pkg="org.mockito" class="AssignmentsManagerTest"/>
<allow pkg="org.opentest4j" class="RemoteLogManagerTest"/>

<!-- see KIP-544 for why KafkaYammerMetrics should be used instead of the global default yammer metrics registry
https://cwiki.apache.org/confluence/display/KAFKA/KIP-544%3A+Make+metrics+exposed+via+JMX+configurable -->
Expand Down
2 changes: 2 additions & 0 deletions checkstyle/suppressions.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@
<suppress checks="ClassFanOutComplexity" files="RemoteLogManagerTest.java"/>
<suppress checks="MethodLength"
files="(KafkaClusterTestKit).java"/>
<suppress checks="JavaNCSS"
files="RemoteLogManagerTest.java"/>

<!-- server tests -->
<suppress checks="MethodLength|JavaNCSS|NPath" files="DescribeTopicPartitionsRequestHandlerTest.java"/>
Expand Down
29 changes: 29 additions & 0 deletions core/src/main/java/kafka/log/remote/RemoteLogManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@
import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.security.PrivilegedAction;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
Expand All @@ -115,6 +116,8 @@
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
Expand Down Expand Up @@ -150,6 +153,8 @@ public class RemoteLogManager implements Closeable {

private final RemoteLogMetadataManager remoteLogMetadataManager;

private final ReentrantLock copyQuotaManagerLock = new ReentrantLock();
private final Condition copyQuotaManagerLockCondition = copyQuotaManagerLock.newCondition();
private final RLMQuotaManager rlmCopyQuotaManager;
private final RLMQuotaManager rlmFetchQuotaManager;

Expand Down Expand Up @@ -234,6 +239,13 @@ private void removeMetrics() {
remoteStorageReaderThreadPool.removeMetrics();
}

/**
* Returns the timeout for the RLM Tasks to wait for the quota to be available
*/
Duration quotaTimeout() {
return Duration.ofSeconds(1);
}

RLMQuotaManager createRLMCopyQuotaManager() {
return new RLMQuotaManager(copyQuotaManagerConfig(rlmConfig), metrics, QuotaType.RLMCopy$.MODULE$,
"Tracking copy byte-rate for Remote Log Manager", time);
Expand Down Expand Up @@ -730,6 +742,23 @@ public void copyLogSegmentsToRemote(UnifiedLog log) throws InterruptedException
isCancelled(), isLeader());
return;
}

copyQuotaManagerLock.lock();
try {
while (rlmCopyQuotaManager.isQuotaExceeded()) {
logger.debug("Quota exceeded for copying log segments, waiting for the quota to be available.");
try {
copyQuotaManagerLockCondition.await(quotaTimeout().toMillis(), TimeUnit.MILLISECONDS);
} catch (InterruptedException e) {
logger.warn("Interrupted while waiting for the quota to be available for copying log segments.");
}
}
rlmCopyQuotaManager.record(candidateLogSegment.logSegment.log().sizeInBytes());
// Signal waiting threads to check the quota again
copyQuotaManagerLockCondition.signalAll();
} finally {
copyQuotaManagerLock.unlock();
}
copyLogSegment(log, candidateLogSegment.logSegment, candidateLogSegment.nextSegmentOffset);
}
}
Expand Down
235 changes: 235 additions & 0 deletions core/src/test/java/kafka/log/remote/RemoteLogManagerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import kafka.cluster.EndPoint;
import kafka.cluster.Partition;
import kafka.log.UnifiedLog;
import kafka.log.remote.quota.RLMQuotaManager;
import kafka.log.remote.quota.RLMQuotaManagerConfig;
import kafka.server.BrokerTopicStats;
import kafka.server.KafkaConfig;
Expand Down Expand Up @@ -85,6 +86,7 @@
import org.mockito.InOrder;
import org.mockito.MockedConstruction;
import org.mockito.Mockito;
import org.opentest4j.AssertionFailedError;
import scala.Option;
import scala.collection.JavaConverters;

Expand All @@ -95,6 +97,7 @@
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
Expand Down Expand Up @@ -134,6 +137,7 @@
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
Expand Down Expand Up @@ -170,9 +174,11 @@ public class RemoteLogManagerTest {
private final String remoteLogMetadataConsumerTestProp = REMOTE_LOG_METADATA_CONSUMER_PREFIX + "consumer.test";
private final String remoteLogMetadataConsumerTestVal = "consumer.test";
private final String remoteLogMetadataTopicPartitionsNum = "1";
private static final String EXPECTED_THE_OPERATION_TO_TIME_OUT = "Expected the operation to time out";

private final RemoteStorageManager remoteStorageManager = mock(RemoteStorageManager.class);
private final RemoteLogMetadataManager remoteLogMetadataManager = mock(RemoteLogMetadataManager.class);
private final RLMQuotaManager rlmCopyQuotaManager = mock(RLMQuotaManager.class);
private RemoteLogManagerConfig remoteLogManagerConfig = null;

private BrokerTopicStats brokerTopicStats = null;
Expand Down Expand Up @@ -225,6 +231,12 @@ public RemoteStorageManager createRemoteStorageManager() {
public RemoteLogMetadataManager createRemoteLogMetadataManager() {
return remoteLogMetadataManager;
}
public RLMQuotaManager createRLMCopyQuotaManager() {
return rlmCopyQuotaManager;
}
public Duration quotaTimeout() {
return Duration.ofMillis(100);
}
@Override
long findLogStartOffset(TopicIdPartition topicIdPartition, UnifiedLog log) {
return 0L;
Expand Down Expand Up @@ -2529,6 +2541,229 @@ public void testFetchQuotaManagerConfig() {
assertEquals(1, rlmFetchQuotaManagerConfig.quotaWindowSizeSeconds());
}

@Test
public void testCopyQuotaNotExceeded() throws Exception {
long oldSegmentStartOffset = 0L;
long nextSegmentStartOffset = 150L;

when(mockLog.topicPartition()).thenReturn(leaderTopicIdPartition.topicPartition());

// leader epoch preparation
checkpoint.write(totalEpochEntries);
LeaderEpochFileCache cache = new LeaderEpochFileCache(leaderTopicIdPartition.topicPartition(), checkpoint);
when(mockLog.leaderEpochCache()).thenReturn(Option.apply(cache));
when(remoteLogMetadataManager.highestOffsetForEpoch(any(TopicIdPartition.class), anyInt())).thenReturn(Optional.of(0L));

// create 2 log segments, with 0 and 150 as log start offset
LogSegment oldSegment = mock(LogSegment.class);
LogSegment activeSegment = mock(LogSegment.class);

File tempFile = TestUtils.tempFile();
FileRecords fileRecords = mock(FileRecords.class);
when(fileRecords.file()).thenReturn(tempFile);
when(fileRecords.sizeInBytes()).thenReturn(10);

// Set up the segment that will be copied
when(oldSegment.log()).thenReturn(fileRecords);
when(oldSegment.baseOffset()).thenReturn(oldSegmentStartOffset);
when(oldSegment.readNextOffset()).thenReturn(nextSegmentStartOffset);

// set up the active segment
when(activeSegment.baseOffset()).thenReturn(nextSegmentStartOffset);

when(mockLog.activeSegment()).thenReturn(activeSegment);
when(mockLog.logStartOffset()).thenReturn(oldSegmentStartOffset);
when(mockLog.logSegments(anyLong(), anyLong())).thenReturn(JavaConverters.collectionAsScalaIterable(Arrays.asList(oldSegment, activeSegment)));

File mockProducerSnapshotIndex = TestUtils.tempFile();
ProducerStateManager mockStateManager = mock(ProducerStateManager.class);
when(mockStateManager.fetchSnapshot(anyLong())).thenReturn(Optional.of(mockProducerSnapshotIndex));

when(mockLog.producerStateManager()).thenReturn(mockStateManager);
when(mockLog.lastStableOffset()).thenReturn(250L);

File tempDir = TestUtils.tempDirectory();
OffsetIndex idx = LazyIndex.forOffset(LogFileUtils.offsetIndexFile(tempDir, oldSegmentStartOffset, ""), oldSegmentStartOffset, 1000).get();
TimeIndex timeIdx = LazyIndex.forTime(LogFileUtils.timeIndexFile(tempDir, oldSegmentStartOffset, ""), oldSegmentStartOffset, 1500).get();
File txnFile = UnifiedLog.transactionIndexFile(tempDir, oldSegmentStartOffset, "");
txnFile.createNewFile();
TransactionIndex txnIndex = new TransactionIndex(oldSegmentStartOffset, txnFile);
when(oldSegment.timeIndex()).thenReturn(timeIdx);
when(oldSegment.offsetIndex()).thenReturn(idx);
when(oldSegment.txnIndex()).thenReturn(txnIndex);

CompletableFuture<Void> dummyFuture = new CompletableFuture<>();
dummyFuture.complete(null);
when(remoteLogMetadataManager.addRemoteLogSegmentMetadata(any(RemoteLogSegmentMetadata.class))).thenReturn(dummyFuture);
when(remoteLogMetadataManager.updateRemoteLogSegmentMetadata(any(RemoteLogSegmentMetadataUpdate.class))).thenReturn(dummyFuture);
when(remoteStorageManager.copyLogSegmentData(any(RemoteLogSegmentMetadata.class), any(LogSegmentData.class))).thenReturn(Optional.empty());

// QuotaManager returns false to indicate quota has not been exceeded
when(rlmCopyQuotaManager.isQuotaExceeded()).thenReturn(false);
doNothing().when(rlmCopyQuotaManager).record(anyInt());

RemoteLogManager.RLMTask task = remoteLogManager.new RLMTask(leaderTopicIdPartition, 128);
task.convertToLeader(2);
// Verify the copy operation completes within the timeout, since it does not need to wait for quota availability
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> task.copyLogSegmentsToRemote(mockLog));

// Verify quota check was performed
verify(rlmCopyQuotaManager, times(1)).isQuotaExceeded();
// Verify bytes to copy was recorded with the quota manager
verify(rlmCopyQuotaManager, times(1)).record(10);

// Verify the highest offset in remote storage is updated
ArgumentCaptor<Long> capture = ArgumentCaptor.forClass(Long.class);
verify(mockLog, times(2)).updateHighestOffsetInRemoteStorage(capture.capture());
List<Long> capturedValues = capture.getAllValues();
// Verify the highest offset in remote storage was -1L before the copy
assertEquals(-1L, capturedValues.get(0).longValue());
// Verify it was updated to 149L after the copy
assertEquals(149L, capturedValues.get(1).longValue());
}

@Test
public void testCopyQuotaExceeded() throws Exception {
long oldSegmentStartOffset = 0L;
long nextSegmentStartOffset = 150L;

when(mockLog.topicPartition()).thenReturn(leaderTopicIdPartition.topicPartition());

// leader epoch preparation
checkpoint.write(totalEpochEntries);
LeaderEpochFileCache cache = new LeaderEpochFileCache(leaderTopicIdPartition.topicPartition(), checkpoint);
when(mockLog.leaderEpochCache()).thenReturn(Option.apply(cache));
when(remoteLogMetadataManager.highestOffsetForEpoch(any(TopicIdPartition.class), anyInt())).thenReturn(Optional.of(0L));

// create 2 log segments, with 0 and 150 as log start offset
LogSegment oldSegment = mock(LogSegment.class);
LogSegment activeSegment = mock(LogSegment.class);

when(oldSegment.baseOffset()).thenReturn(oldSegmentStartOffset);
when(oldSegment.readNextOffset()).thenReturn(nextSegmentStartOffset);
when(activeSegment.baseOffset()).thenReturn(nextSegmentStartOffset);

when(mockLog.activeSegment()).thenReturn(activeSegment);
when(mockLog.logStartOffset()).thenReturn(oldSegmentStartOffset);
when(mockLog.logSegments(anyLong(), anyLong())).thenReturn(JavaConverters.collectionAsScalaIterable(Arrays.asList(oldSegment, activeSegment)));
when(mockLog.lastStableOffset()).thenReturn(250L);

// QuotaManager returns true to indicate that the quota has been exceeded
when(rlmCopyQuotaManager.isQuotaExceeded()).thenReturn(true);
doNothing().when(rlmCopyQuotaManager).record(anyInt());

RemoteLogManager.RLMTask task = remoteLogManager.new RLMTask(leaderTopicIdPartition, 128);
task.convertToLeader(2);

// Verify that the copy operation times out, since no segments can be copied due to quota being exceeded
try {
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> task.copyLogSegmentsToRemote(mockLog));
fail(EXPECTED_THE_OPERATION_TO_TIME_OUT);
} catch (AssertionFailedError e) {
// Fail the test if the operation completed within the timeout
if (e.getMessage().equals(EXPECTED_THE_OPERATION_TO_TIME_OUT)) {
fail(e.getMessage());
}
}

// Verify the highest offset in remote storage is updated only once
ArgumentCaptor<Long> capture = ArgumentCaptor.forClass(Long.class);
verify(mockLog, times(1)).updateHighestOffsetInRemoteStorage(capture.capture());
// Verify the highest offset in remote storage was -1L before the copy started
assertEquals(-1L, capture.getValue());
}

@Test
public void testCopyThrottling() throws Exception {
long oldestSegmentStartOffset = 0L;

when(mockLog.topicPartition()).thenReturn(leaderTopicIdPartition.topicPartition());

// leader epoch preparation
checkpoint.write(totalEpochEntries);
LeaderEpochFileCache cache = new LeaderEpochFileCache(leaderTopicIdPartition.topicPartition(), checkpoint);
when(mockLog.leaderEpochCache()).thenReturn(Option.apply(cache));
when(remoteLogMetadataManager.highestOffsetForEpoch(any(TopicIdPartition.class), anyInt())).thenReturn(Optional.of(0L));

// create 3 log segments
LogSegment segmentToCopy = mock(LogSegment.class);
LogSegment segmentToThrottle = mock(LogSegment.class);
LogSegment activeSegment = mock(LogSegment.class);

File tempFile = TestUtils.tempFile();
FileRecords fileRecords = mock(FileRecords.class);
when(fileRecords.file()).thenReturn(tempFile);
when(fileRecords.sizeInBytes()).thenReturn(10);

// set up the segment that will be copied
when(segmentToCopy.log()).thenReturn(fileRecords);
when(segmentToCopy.baseOffset()).thenReturn(oldestSegmentStartOffset);
when(segmentToCopy.readNextOffset()).thenReturn(100L);

// set up the segment that will not be copied because of hitting quota
when(segmentToThrottle.log()).thenReturn(fileRecords);
when(segmentToThrottle.baseOffset()).thenReturn(100L);
when(segmentToThrottle.readNextOffset()).thenReturn(150L);

// set up the active segment
when(activeSegment.log()).thenReturn(fileRecords);
when(activeSegment.baseOffset()).thenReturn(150L);

when(mockLog.activeSegment()).thenReturn(activeSegment);
when(mockLog.logStartOffset()).thenReturn(oldestSegmentStartOffset);
when(mockLog.logSegments(anyLong(), anyLong())).thenReturn(JavaConverters.collectionAsScalaIterable(Arrays.asList(segmentToCopy, segmentToThrottle, activeSegment)));

File mockProducerSnapshotIndex = TestUtils.tempFile();
ProducerStateManager mockStateManager = mock(ProducerStateManager.class);
when(mockStateManager.fetchSnapshot(anyLong())).thenReturn(Optional.of(mockProducerSnapshotIndex));

when(mockLog.producerStateManager()).thenReturn(mockStateManager);
when(mockLog.lastStableOffset()).thenReturn(250L);

File tempDir = TestUtils.tempDirectory();
OffsetIndex idx = LazyIndex.forOffset(LogFileUtils.offsetIndexFile(tempDir, oldestSegmentStartOffset, ""), oldestSegmentStartOffset, 1000).get();
TimeIndex timeIdx = LazyIndex.forTime(LogFileUtils.timeIndexFile(tempDir, oldestSegmentStartOffset, ""), oldestSegmentStartOffset, 1500).get();
File txnFile = UnifiedLog.transactionIndexFile(tempDir, oldestSegmentStartOffset, "");
txnFile.createNewFile();
TransactionIndex txnIndex = new TransactionIndex(oldestSegmentStartOffset, txnFile);
when(segmentToCopy.timeIndex()).thenReturn(timeIdx);
when(segmentToCopy.offsetIndex()).thenReturn(idx);
when(segmentToCopy.txnIndex()).thenReturn(txnIndex);

CompletableFuture<Void> dummyFuture = new CompletableFuture<>();
dummyFuture.complete(null);
when(remoteLogMetadataManager.addRemoteLogSegmentMetadata(any(RemoteLogSegmentMetadata.class))).thenReturn(dummyFuture);
when(remoteLogMetadataManager.updateRemoteLogSegmentMetadata(any(RemoteLogSegmentMetadataUpdate.class))).thenReturn(dummyFuture);
when(remoteStorageManager.copyLogSegmentData(any(RemoteLogSegmentMetadata.class), any(LogSegmentData.class))).thenReturn(Optional.empty());

// After the first call, isQuotaExceeded should return true
when(rlmCopyQuotaManager.isQuotaExceeded()).thenReturn(false, true);
doNothing().when(rlmCopyQuotaManager).record(anyInt());

RemoteLogManager.RLMTask task = remoteLogManager.new RLMTask(leaderTopicIdPartition, 128);
task.convertToLeader(2);

// Verify that the copy operation times out, since the second segment cannot be copied due to quota being exceeded
try {
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> task.copyLogSegmentsToRemote(mockLog));
fail(EXPECTED_THE_OPERATION_TO_TIME_OUT);
} catch (AssertionFailedError e) {
// Fail the test if the operation completed within the timeout
if (e.getMessage().equals(EXPECTED_THE_OPERATION_TO_TIME_OUT)) {
fail(e.getMessage());
}
}

// Verify the highest offset in remote storage is updated
ArgumentCaptor<Long> capture = ArgumentCaptor.forClass(Long.class);
verify(mockLog, times(2)).updateHighestOffsetInRemoteStorage(capture.capture());
List<Long> capturedValues = capture.getAllValues();
// Verify the highest offset in remote storage was -1L before the copy
assertEquals(-1L, capturedValues.get(0).longValue());
// Verify it was updated to 99L after the copy
assertEquals(99L, capturedValues.get(1).longValue());
}

private Partition mockPartition(TopicIdPartition topicIdPartition) {
TopicPartition tp = topicIdPartition.topicPartition();
Partition partition = mock(Partition.class);
Expand Down

0 comments on commit 3593def

Please sign in to comment.