Skip to content

Commit

Permalink
Fix: recover command didn't accept rate limit parameter
Browse files Browse the repository at this point in the history
  • Loading branch information
@dlg99
dlg99 committed Dec 10, 2024
1 parent b372e90 commit 5156ac6
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 31 deletions.
2 changes: 1 addition & 1 deletion bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/BookieShell.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -524,7 +524,7 @@ public RecoverCmd() {
opts.addOption("sk", "skipOpenLedgers", false, "Skip recovering open ledgers");
opts.addOption("d", "deleteCookie", false, "Delete cookie node for the bookie.");
opts.addOption("sku", "skipUnrecoverableLedgers", false, "Skip unrecoverable ledgers.");
opts.addOption("rate", "replicationRate", false, "Replication rate by bytes");
opts.addOption("rate", "replicationRate", true, "Replication rate by bytes");
}

@Override
Expand Down
98 changes: 68 additions & 30 deletions bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/BookieShellTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
Expand All @@ -35,6 +36,9 @@
import static org.mockito.Mockito.when;

import com.google.common.collect.Maps;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.util.Set;
import java.util.SortedMap;
import java.util.function.Function;
Expand Down Expand Up @@ -189,6 +193,14 @@ public void testRecoverCmdQuery() throws Exception {
verify(admin, times(1)).close();
}

@SuppressWarnings("unchecked")
@Test
public void testRecoverLedgerWithRateLimit() throws Exception {
testRecoverCmdRecoverLedger(
12345, false, false, false,
"-force", "-l", "12345", "-rate", "10000", "127.0.0.1:3181");
}

@Test
public void testRecoverCmdRecoverLedgerDefault() throws Exception {
// default behavior
Expand Down Expand Up @@ -235,21 +247,30 @@ void testRecoverCmdRecoverLedger(long ledgerId,
boolean skipOpenLedgers,
boolean removeCookies,
String... args) throws Exception {
RecoverCmd cmd = (RecoverCmd) shell.commands.get("recover");
CommandLine cmdLine = parseCommandLine(cmd, args);
assertEquals(0, cmd.runCmd(cmdLine));
bookKeeperAdminMockedStatic.verify(() -> BookKeeperAdmin.newBookKeeperAdmin(any(ClientConfiguration.class)),
times(1));
verify(admin, times(1))
.recoverBookieData(eq(ledgerId), any(Set.class), eq(dryrun), eq(skipOpenLedgers));
verify(admin, times(1)).close();
if (removeCookies) {
MetadataDrivers.runFunctionWithRegistrationManager(any(ServerConfiguration.class), any(Function.class));
verify(rm, times(1)).readCookie(any(BookieId.class));
verify(rm, times(1)).removeCookie(any(BookieId.class), eq(version));
} else {
verify(rm, times(0)).readCookie(any(BookieId.class));
verify(rm, times(0)).removeCookie(any(BookieId.class), eq(version));
final PrintStream oldPs = System.err;
try(ByteArrayOutputStream outContent = new ByteArrayOutputStream()) {
System.setErr(new PrintStream(outContent));

RecoverCmd cmd = (RecoverCmd) shell.commands.get("recover");
CommandLine cmdLine = parseCommandLine(cmd, args);
assertEquals(0, cmd.runCmd(cmdLine));
bookKeeperAdminMockedStatic.verify(() -> BookKeeperAdmin.newBookKeeperAdmin(any(ClientConfiguration.class)),
times(1));
verify(admin, times(1))
.recoverBookieData(eq(ledgerId), any(Set.class), eq(dryrun), eq(skipOpenLedgers));
verify(admin, times(1)).close();
if (removeCookies) {
MetadataDrivers.runFunctionWithRegistrationManager(any(ServerConfiguration.class), any(Function.class));
verify(rm, times(1)).readCookie(any(BookieId.class));
verify(rm, times(1)).removeCookie(any(BookieId.class), eq(version));
} else {
verify(rm, times(0)).readCookie(any(BookieId.class));
verify(rm, times(0)).removeCookie(any(BookieId.class), eq(version));
}
assertFalse("invalid value for option detected:\n" + outContent,
outContent.toString().contains("invalid value for option"));
} finally {
System.setErr(oldPs);
}
}

Expand All @@ -261,6 +282,14 @@ public void testRecoverCmdRecoverDefault() throws Exception {
"-force", "127.0.0.1:3181");
}

@Test
public void testRecoverWithRateLimit() throws Exception {
// default behavior
testRecoverCmdRecover(
false, false, false, false,
"-force", "127.0.0.1:3181");
}

@Test
public void testRecoverCmdRecoverDeleteCookie() throws Exception {
// dryrun
Expand Down Expand Up @@ -307,21 +336,30 @@ void testRecoverCmdRecover(boolean dryrun,
boolean removeCookies,
boolean skipUnrecoverableLedgers,
String... args) throws Exception {
RecoverCmd cmd = (RecoverCmd) shell.commands.get("recover");
CommandLine cmdLine = parseCommandLine(cmd, args);
assertEquals(0, cmd.runCmd(cmdLine));
bookKeeperAdminMockedStatic.verify(() -> BookKeeperAdmin.newBookKeeperAdmin(any(ClientConfiguration.class)),
times(1));
verify(admin, times(1))
.recoverBookieData(any(Set.class), eq(dryrun), eq(skipOpenLedgers), eq(skipUnrecoverableLedgers));
verify(admin, times(1)).close();
if (removeCookies) {
MetadataDrivers.runFunctionWithRegistrationManager(any(ServerConfiguration.class), any(Function.class));
verify(rm, times(1)).readCookie(any(BookieId.class));
verify(rm, times(1)).removeCookie(any(BookieId.class), eq(version));
} else {
verify(rm, times(0)).readCookie(any(BookieId.class));
verify(rm, times(0)).removeCookie(any(BookieId.class), eq(version));
final PrintStream oldPs = System.err;
try(ByteArrayOutputStream outContent = new ByteArrayOutputStream()) {
System.setErr(new PrintStream(outContent));

RecoverCmd cmd = (RecoverCmd) shell.commands.get("recover");
CommandLine cmdLine = parseCommandLine(cmd, args);
assertEquals(0, cmd.runCmd(cmdLine));
bookKeeperAdminMockedStatic.verify(() -> BookKeeperAdmin.newBookKeeperAdmin(any(ClientConfiguration.class)),
times(1));
verify(admin, times(1))
.recoverBookieData(any(Set.class), eq(dryrun), eq(skipOpenLedgers), eq(skipUnrecoverableLedgers));
verify(admin, times(1)).close();
if (removeCookies) {
MetadataDrivers.runFunctionWithRegistrationManager(any(ServerConfiguration.class), any(Function.class));
verify(rm, times(1)).readCookie(any(BookieId.class));
verify(rm, times(1)).removeCookie(any(BookieId.class), eq(version));
} else {
verify(rm, times(0)).readCookie(any(BookieId.class));
verify(rm, times(0)).removeCookie(any(BookieId.class), eq(version));
}
assertFalse("invalid value for option detected:\n" + outContent,
outContent.toString().contains("invalid value for option"));
} finally {
System.setErr(oldPs);
}
}

Expand Down

0 comments on commit 5156ac6

Please sign in to comment.