| /** |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.bookkeeper.client; |
| |
| import static com.google.common.base.Preconditions.checkState; |
| import static org.apache.bookkeeper.client.api.BKException.Code.NoBookieAvailableException; |
| import static org.mockito.ArgumentMatchers.any; |
| import static org.mockito.ArgumentMatchers.anyBoolean; |
| import static org.mockito.ArgumentMatchers.anyInt; |
| import static org.mockito.ArgumentMatchers.anyList; |
| import static org.mockito.ArgumentMatchers.anyLong; |
| import static org.mockito.ArgumentMatchers.anyMap; |
| import static org.mockito.ArgumentMatchers.anySet; |
| import static org.mockito.Mockito.doAnswer; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.when; |
| |
| import io.netty.buffer.ByteBuf; |
| import io.netty.buffer.ByteBufAllocator; |
| import io.netty.buffer.Unpooled; |
| import io.netty.buffer.UnpooledByteBufAllocator; |
| |
| import java.security.GeneralSecurityException; |
| import java.util.ArrayList; |
| import java.util.Collections; |
| import java.util.EnumSet; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.concurrent.CompletableFuture; |
| import java.util.concurrent.ConcurrentHashMap; |
| import java.util.concurrent.ConcurrentMap; |
| import java.util.concurrent.ConcurrentSkipListSet; |
| import java.util.concurrent.CopyOnWriteArrayList; |
| import java.util.concurrent.atomic.AtomicLong; |
| import java.util.concurrent.locks.ReentrantReadWriteLock; |
| |
| import org.apache.bookkeeper.client.BKException.BKDigestMatchException; |
| import org.apache.bookkeeper.client.BKException.Code; |
| import org.apache.bookkeeper.client.api.CreateBuilder; |
| import org.apache.bookkeeper.client.api.DeleteBuilder; |
| import org.apache.bookkeeper.client.api.LedgerMetadata; |
| import org.apache.bookkeeper.client.api.OpenBuilder; |
| import org.apache.bookkeeper.common.util.OrderedExecutor; |
| import org.apache.bookkeeper.common.util.OrderedScheduler; |
| import org.apache.bookkeeper.conf.ClientConfiguration; |
| import org.apache.bookkeeper.meta.LedgerIdGenerator; |
| import org.apache.bookkeeper.meta.LedgerManager; |
| import org.apache.bookkeeper.net.BookieSocketAddress; |
| import org.apache.bookkeeper.proto.BookieClient; |
| import org.apache.bookkeeper.proto.BookieProtocol; |
| import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks; |
| import org.apache.bookkeeper.proto.checksum.DigestManager; |
| import org.apache.bookkeeper.stats.NullStatsLogger; |
| import org.apache.bookkeeper.util.ByteBufList; |
| import org.apache.bookkeeper.versioning.LongVersion; |
| import org.apache.bookkeeper.versioning.Version; |
| import org.apache.bookkeeper.versioning.Versioned; |
| import org.junit.After; |
| import org.junit.Before; |
| import org.mockito.invocation.InvocationOnMock; |
| import org.mockito.stubbing.Answer; |
| import org.mockito.stubbing.Stubber; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Base class for Mock-based Client testcases. |
| */ |
| public abstract class MockBookKeeperTestCase { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(MockBookKeeperTestCase.class); |
| |
| protected OrderedScheduler scheduler; |
| protected OrderedExecutor executor; |
| protected BookKeeper bk; |
| protected BookieClient bookieClient; |
| protected LedgerManager ledgerManager; |
| protected LedgerIdGenerator ledgerIdGenerator; |
| |
| private BookieWatcher bookieWatcher; |
| |
| protected ConcurrentMap<Long, LedgerMetadata> mockLedgerMetadataRegistry; |
| protected AtomicLong mockNextLedgerId; |
| protected ConcurrentSkipListSet<Long> fencedLedgers; |
| protected ConcurrentMap<Long, Map<BookieSocketAddress, Map<Long, MockEntry>>> mockLedgerData; |
| |
| private Map<BookieSocketAddress, List<Runnable>> deferredBookieForceLedgerResponses; |
| private Set<BookieSocketAddress> suspendedBookiesForForceLedgerAcks; |
| |
| List<BookieSocketAddress> failedBookies; |
| Set<BookieSocketAddress> availableBookies; |
| private int lastIndexForBK; |
| |
| private Map<BookieSocketAddress, Map<Long, MockEntry>> getMockLedgerContents(long ledgerId) { |
| return mockLedgerData.computeIfAbsent(ledgerId, (id) -> new ConcurrentHashMap<>()); |
| } |
| |
| private Map<Long, MockEntry> getMockLedgerContentsInBookie(long ledgerId, BookieSocketAddress bookieSocketAddress) { |
| return getMockLedgerContents(ledgerId).computeIfAbsent(bookieSocketAddress, addr -> new ConcurrentHashMap<>()); |
| } |
| |
| private MockEntry getMockLedgerEntry(long ledgerId, |
| BookieSocketAddress bookieSocketAddress, long entryId) throws BKException{ |
| if (failedBookies.contains(bookieSocketAddress)) { |
| throw BKException.create(NoBookieAvailableException); |
| } |
| return getMockLedgerContentsInBookie(ledgerId, bookieSocketAddress).get(entryId); |
| } |
| |
| private static final class MockEntry { |
| |
| byte[] payload; |
| long lastAddConfirmed; |
| |
| public MockEntry(byte[] payload, long lastAddConfirmed) { |
| this.payload = payload; |
| this.lastAddConfirmed = lastAddConfirmed; |
| } |
| |
| } |
| |
| @Before |
| public void setup() throws Exception { |
| deferredBookieForceLedgerResponses = new ConcurrentHashMap<>(); |
| suspendedBookiesForForceLedgerAcks = Collections.synchronizedSet(new HashSet<>()); |
| mockLedgerMetadataRegistry = new ConcurrentHashMap<>(); |
| mockLedgerData = new ConcurrentHashMap<>(); |
| mockNextLedgerId = new AtomicLong(1); |
| fencedLedgers = new ConcurrentSkipListSet<>(); |
| scheduler = OrderedScheduler.newSchedulerBuilder().numThreads(4).name("bk-test").build(); |
| executor = OrderedExecutor.newBuilder().build(); |
| bookieWatcher = mock(BookieWatcher.class); |
| |
| bookieClient = mock(BookieClient.class); |
| ledgerManager = mock(LedgerManager.class); |
| ledgerIdGenerator = mock(LedgerIdGenerator.class); |
| |
| bk = mock(BookKeeper.class); |
| |
| failedBookies = new ArrayList<>(); |
| availableBookies = new HashSet<>(); |
| |
| when(bk.getCloseLock()).thenReturn(new ReentrantReadWriteLock()); |
| when(bk.isClosed()).thenReturn(false); |
| when(bk.getBookieWatcher()).thenReturn(bookieWatcher); |
| when(bk.getMainWorkerPool()).thenReturn(executor); |
| when(bk.getBookieClient()).thenReturn(bookieClient); |
| when(bk.getScheduler()).thenReturn(scheduler); |
| |
| setBookKeeperConfig(new ClientConfiguration()); |
| when(bk.getStatsLogger()).thenReturn(NullStatsLogger.INSTANCE); |
| BookKeeperClientStats clientStats = BookKeeperClientStats.newInstance(NullStatsLogger.INSTANCE); |
| ClientContext clientCtx = new ClientContext() { |
| @Override |
| public ClientInternalConf getConf() { |
| return ClientInternalConf.fromConfig(bk.getConf()); |
| } |
| |
| @Override |
| public LedgerManager getLedgerManager() { |
| return ledgerManager; |
| } |
| |
| @Override |
| public BookieWatcher getBookieWatcher() { |
| return bookieWatcher; |
| } |
| |
| @Override |
| public EnsemblePlacementPolicy getPlacementPolicy() { |
| return null; |
| } |
| |
| @Override |
| public BookieClient getBookieClient() { |
| return bookieClient; |
| } |
| |
| @Override |
| public OrderedExecutor getMainWorkerPool() { |
| return scheduler; |
| } |
| |
| @Override |
| public OrderedScheduler getScheduler() { |
| return scheduler; |
| } |
| |
| @Override |
| public BookKeeperClientStats getClientStats() { |
| return clientStats; |
| } |
| |
| @Override |
| public boolean isClientClosed() { |
| return bk.isClosed(); |
| } |
| |
| @Override |
| public ByteBufAllocator getByteBufAllocator() { |
| return UnpooledByteBufAllocator.DEFAULT; |
| } |
| }; |
| when(bk.getClientCtx()).thenReturn(clientCtx); |
| when(bk.getLedgerManager()).thenReturn(ledgerManager); |
| when(bk.getLedgerIdGenerator()).thenReturn(ledgerIdGenerator); |
| when(bk.getReturnRc(anyInt())).thenAnswer(invocationOnMock -> invocationOnMock.getArgument(0)); |
| when(bookieClient.isWritable(any(), anyLong())).thenReturn(true); |
| |
| setupLedgerIdGenerator(); |
| setupCreateLedgerMetadata(); |
| setupReadLedgerMetadata(); |
| setupWriteLedgerMetadata(); |
| setupRemoveLedgerMetadata(); |
| setupRegisterLedgerMetadataListener(); |
| setupBookieWatcherForNewEnsemble(); |
| setupBookieWatcherForEnsembleChange(); |
| setupBookieClientReadEntry(); |
| setupBookieClientAddEntry(); |
| setupBookieClientForceLedger(); |
| } |
| |
| protected void setBookKeeperConfig(ClientConfiguration conf) { |
| when(bk.getConf()).thenReturn(conf); |
| } |
| |
| private DigestManager getDigestType(long ledgerId) throws GeneralSecurityException { |
| LedgerMetadata metadata = mockLedgerMetadataRegistry.get(ledgerId); |
| return DigestManager.instantiate( |
| ledgerId, |
| metadata.getPassword(), |
| org.apache.bookkeeper.client.BookKeeper.DigestType.toProtoDigestType( |
| org.apache.bookkeeper.client.BookKeeper.DigestType.fromApiDigestType( |
| metadata.getDigestType())), |
| UnpooledByteBufAllocator.DEFAULT, false); |
| } |
| |
| @After |
| public void tearDown() { |
| scheduler.shutdown(); |
| executor.shutdown(); |
| } |
| |
| protected CreateBuilder newCreateLedgerOp() { |
| return new LedgerCreateOp.CreateBuilderImpl(bk); |
| } |
| |
| protected OpenBuilder newOpenLedgerOp() { |
| return new LedgerOpenOp.OpenBuilderImpl(bk); |
| } |
| |
| protected DeleteBuilder newDeleteLedgerOp() { |
| return new LedgerDeleteOp.DeleteBuilderImpl(bk); |
| } |
| |
| protected void closeBookkeeper() { |
| when(bk.isClosed()).thenReturn(true); |
| } |
| |
| protected void killBookie(BookieSocketAddress killedBookieSocketAddress) { |
| failedBookies.add(killedBookieSocketAddress); |
| availableBookies.remove(killedBookieSocketAddress); |
| } |
| |
| protected void startKilledBookie(BookieSocketAddress killedBookieSocketAddress) { |
| checkState(failedBookies.contains(killedBookieSocketAddress)); |
| checkState(!availableBookies.contains(killedBookieSocketAddress)); |
| failedBookies.remove(killedBookieSocketAddress); |
| availableBookies.add(killedBookieSocketAddress); |
| } |
| |
| protected void suspendBookieForceLedgerAcks(BookieSocketAddress address) { |
| suspendedBookiesForForceLedgerAcks.add(address); |
| } |
| |
| protected void resumeBookieWriteAcks(BookieSocketAddress address) { |
| suspendedBookiesForForceLedgerAcks.remove(address); |
| List<Runnable> pendingResponses = deferredBookieForceLedgerResponses.remove(address); |
| if (pendingResponses != null) { |
| pendingResponses.forEach(Runnable::run); |
| } |
| } |
| |
| protected BookieSocketAddress startNewBookie() { |
| BookieSocketAddress address = generateBookieSocketAddress(lastIndexForBK++); |
| availableBookies.add(address); |
| return address; |
| } |
| |
| protected BookieSocketAddress generateBookieSocketAddress(int index) { |
| return new BookieSocketAddress("localhost", 1111 + index); |
| } |
| |
| protected ArrayList<BookieSocketAddress> generateNewEnsemble(int ensembleSize) { |
| ArrayList<BookieSocketAddress> ensemble = new ArrayList<>(ensembleSize); |
| for (int i = 0; i < ensembleSize; i++) { |
| ensemble.add(generateBookieSocketAddress(i)); |
| } |
| availableBookies.addAll(ensemble); |
| lastIndexForBK = ensembleSize; |
| return ensemble; |
| } |
| |
| private void setupBookieWatcherForNewEnsemble() throws BKException.BKNotEnoughBookiesException { |
| when(bookieWatcher.newEnsemble(anyInt(), anyInt(), anyInt(), any())) |
| .thenAnswer((Answer<ArrayList<BookieSocketAddress>>) new Answer<ArrayList<BookieSocketAddress>>() { |
| @Override |
| @SuppressWarnings("unchecked") |
| public ArrayList<BookieSocketAddress> answer(InvocationOnMock invocation) throws Throwable { |
| Object[] args = invocation.getArguments(); |
| int ensembleSize = (Integer) args[0]; |
| return generateNewEnsemble(ensembleSize); |
| } |
| }); |
| } |
| |
| private void setupBookieWatcherForEnsembleChange() throws BKException.BKNotEnoughBookiesException { |
| when(bookieWatcher.replaceBookie(anyInt(), anyInt(), anyInt(), anyMap(), anyList(), anyInt(), anySet())) |
| .thenAnswer((Answer<BookieSocketAddress>) new Answer<BookieSocketAddress>() { |
| @Override |
| @SuppressWarnings("unchecked") |
| public BookieSocketAddress answer(InvocationOnMock invocation) throws Throwable { |
| Object[] args = invocation.getArguments(); |
| List<BookieSocketAddress> existingBookies = (List<BookieSocketAddress>) args[4]; |
| Set<BookieSocketAddress> excludeBookies = (Set<BookieSocketAddress>) args[6]; |
| excludeBookies.addAll(existingBookies); |
| Set<BookieSocketAddress> remainBookies = new HashSet<BookieSocketAddress>(availableBookies); |
| remainBookies.removeAll(excludeBookies); |
| if (remainBookies.iterator().hasNext()) { |
| return remainBookies.iterator().next(); |
| } |
| throw BKException.create(BKException.Code.NotEnoughBookiesException); |
| } |
| }); |
| } |
| |
| protected void registerMockEntryForRead(long ledgerId, long entryId, BookieSocketAddress bookieSocketAddress, |
| byte[] entryData, long lastAddConfirmed) { |
| getMockLedgerContentsInBookie(ledgerId, bookieSocketAddress).put(entryId, new MockEntry(entryData, |
| lastAddConfirmed)); |
| } |
| |
| protected void registerMockLedgerMetadata(long ledgerId, LedgerMetadata ledgerMetadata) { |
| mockLedgerMetadataRegistry.put(ledgerId, ledgerMetadata); |
| } |
| |
| protected void setNewGeneratedLedgerId(long ledgerId) { |
| mockNextLedgerId.set(ledgerId); |
| setupLedgerIdGenerator(); |
| } |
| |
| protected LedgerMetadata getLedgerMetadata(long ledgerId) { |
| return mockLedgerMetadataRegistry.get(ledgerId); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void setupReadLedgerMetadata() { |
| doAnswer(invocation -> { |
| Object[] args = invocation.getArguments(); |
| Long ledgerId = (Long) args[0]; |
| CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>(); |
| executor.executeOrdered(ledgerId, () -> { |
| LedgerMetadata ledgerMetadata = mockLedgerMetadataRegistry.get(ledgerId); |
| if (ledgerMetadata == null) { |
| promise.completeExceptionally(new BKException.BKNoSuchLedgerExistsOnMetadataServerException()); |
| } else { |
| promise.complete(new Versioned<>(ledgerMetadata, new LongVersion(1))); |
| } |
| }); |
| return promise; |
| }).when(ledgerManager).readLedgerMetadata(anyLong()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void setupRemoveLedgerMetadata() { |
| doAnswer(invocation -> { |
| Object[] args = invocation.getArguments(); |
| Long ledgerId = (Long) args[0]; |
| CompletableFuture<Void> promise = new CompletableFuture<>(); |
| executor.executeOrdered(ledgerId, () -> { |
| if (mockLedgerMetadataRegistry.remove(ledgerId) != null) { |
| promise.complete(null); |
| } else { |
| promise.completeExceptionally(new BKException.BKNoSuchLedgerExistsOnMetadataServerException()); |
| } |
| }); |
| return promise; |
| }).when(ledgerManager).removeLedgerMetadata(anyLong(), any()); |
| } |
| |
| private void setupRegisterLedgerMetadataListener() { |
| doAnswer((Answer<Void>) new Answer<Void>() { |
| @Override |
| @SuppressWarnings("unchecked") |
| public Void answer(InvocationOnMock invocation) throws Throwable { |
| return null; |
| } |
| }).when(ledgerManager).registerLedgerMetadataListener(anyLong(), any()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void setupLedgerIdGenerator() { |
| doAnswer(invocation -> { |
| Object[] args = invocation.getArguments(); |
| BookkeeperInternalCallbacks.GenericCallback cb = (BookkeeperInternalCallbacks.GenericCallback) args[0]; |
| cb.operationComplete(Code.OK, mockNextLedgerId.getAndIncrement()); |
| return null; |
| }).when(ledgerIdGenerator).generateLedgerId(any()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void setupCreateLedgerMetadata() { |
| doAnswer(invocation -> { |
| Object[] args = invocation.getArguments(); |
| Long ledgerId = (Long) args[0]; |
| |
| CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>(); |
| executor.executeOrdered(ledgerId, () -> { |
| |
| LedgerMetadata ledgerMetadata = (LedgerMetadata) args[1]; |
| mockLedgerMetadataRegistry.put(ledgerId, ledgerMetadata); |
| promise.complete(new Versioned<>(ledgerMetadata, new LongVersion(1))); |
| }); |
| return promise; |
| }).when(ledgerManager).createLedgerMetadata(anyLong(), any()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void setupWriteLedgerMetadata() { |
| doAnswer(invocation -> { |
| Object[] args = invocation.getArguments(); |
| Long ledgerId = (Long) args[0]; |
| LedgerMetadata metadata = (LedgerMetadata) args[1]; |
| Version currentVersion = (Version) args[2]; |
| CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>(); |
| executor.executeOrdered(ledgerId, () -> { |
| LedgerMetadata newMetadata = LedgerMetadataBuilder.from(metadata).build(); |
| mockLedgerMetadataRegistry.put(ledgerId, newMetadata); |
| promise.complete(new Versioned<>(newMetadata, new LongVersion(1234))); |
| }); |
| return promise; |
| }).when(ledgerManager).writeLedgerMetadata(anyLong(), any(), any()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| protected void setupBookieClientReadEntry() { |
| final Stubber stub = doAnswer(invokation -> { |
| Object[] args = invokation.getArguments(); |
| BookieSocketAddress bookieSocketAddress = (BookieSocketAddress) args[0]; |
| long ledgerId = (Long) args[1]; |
| long entryId = (Long) args[2]; |
| BookkeeperInternalCallbacks.ReadEntryCallback callback = |
| (BookkeeperInternalCallbacks.ReadEntryCallback) args[3]; |
| boolean fenced = (((Integer) args[5]) & BookieProtocol.FLAG_DO_FENCING) == BookieProtocol.FLAG_DO_FENCING; |
| |
| executor.executeOrdered(ledgerId, () -> { |
| DigestManager macManager = null; |
| try { |
| macManager = getDigestType(ledgerId); |
| } catch (GeneralSecurityException gse){ |
| LOG.error("Initialize macManager fail", gse); |
| } |
| MockEntry mockEntry = null; |
| try { |
| mockEntry = getMockLedgerEntry(ledgerId, bookieSocketAddress, entryId); |
| } catch (BKException bke) { |
| LOG.info("readEntryAndFenceLedger - occur BKException {}@{} at {}", entryId, ledgerId, |
| bookieSocketAddress); |
| callback.readEntryComplete(bke.getCode(), ledgerId, entryId, null, args[5]); |
| } |
| |
| if (fenced) { |
| fencedLedgers.add(ledgerId); |
| } |
| |
| if (mockEntry != null) { |
| LOG.info("readEntry - found mock entry {}@{} at {}", entryId, ledgerId, bookieSocketAddress); |
| ByteBufList entry = macManager.computeDigestAndPackageForSending(entryId, |
| mockEntry.lastAddConfirmed, mockEntry.payload.length, |
| Unpooled.wrappedBuffer(mockEntry.payload)); |
| callback.readEntryComplete(BKException.Code.OK, ledgerId, entryId, ByteBufList.coalesce(entry), |
| args[4]); |
| entry.release(); |
| } else { |
| LOG.info("readEntry - no such mock entry {}@{} at {}", entryId, ledgerId, bookieSocketAddress); |
| callback.readEntryComplete(BKException.Code.NoSuchEntryException, ledgerId, entryId, null, args[4]); |
| } |
| }); |
| return null; |
| }); |
| |
| stub.when(bookieClient).readEntry(any(), anyLong(), anyLong(), |
| any(BookkeeperInternalCallbacks.ReadEntryCallback.class), |
| any(), anyInt()); |
| |
| stub.when(bookieClient).readEntry(any(), anyLong(), anyLong(), |
| any(BookkeeperInternalCallbacks.ReadEntryCallback.class), |
| any(), anyInt(), any()); |
| |
| stub.when(bookieClient).readEntry(any(), anyLong(), anyLong(), |
| any(BookkeeperInternalCallbacks.ReadEntryCallback.class), |
| any(), anyInt(), any(), anyBoolean()); |
| } |
| |
| private byte[] extractEntryPayload(long ledgerId, long entryId, ByteBufList toSend) |
| throws BKException.BKDigestMatchException { |
| ByteBuf toSendCopy = Unpooled.copiedBuffer(toSend.toArray()); |
| toSendCopy.resetReaderIndex(); |
| DigestManager macManager = null; |
| try { |
| macManager = getDigestType(ledgerId); |
| } catch (GeneralSecurityException gse){ |
| LOG.error("Initialize macManager fail", gse); |
| } |
| ByteBuf content = macManager.verifyDigestAndReturnData(entryId, toSendCopy); |
| byte[] entry = new byte[content.readableBytes()]; |
| content.readBytes(entry); |
| content.resetReaderIndex(); |
| content.release(); |
| return entry; |
| } |
| |
| @SuppressWarnings("unchecked") |
| protected void setupBookieClientAddEntry() { |
| final Stubber stub = doAnswer(invokation -> { |
| Object[] args = invokation.getArguments(); |
| BookkeeperInternalCallbacks.WriteCallback callback = (BookkeeperInternalCallbacks.WriteCallback) args[5]; |
| BookieSocketAddress bookieSocketAddress = (BookieSocketAddress) args[0]; |
| long ledgerId = (Long) args[1]; |
| long entryId = (Long) args[3]; |
| ByteBufList toSend = (ByteBufList) args[4]; |
| Object ctx = args[6]; |
| int options = (int) args[7]; |
| boolean isRecoveryAdd = |
| ((short) options & BookieProtocol.FLAG_RECOVERY_ADD) == BookieProtocol.FLAG_RECOVERY_ADD; |
| |
| toSend.retain(); |
| executor.executeOrdered(ledgerId, () -> { |
| byte[] entry; |
| try { |
| entry = extractEntryPayload(ledgerId, entryId, toSend); |
| } catch (BKDigestMatchException e) { |
| callback.writeComplete(Code.DigestMatchException, |
| ledgerId, entryId, bookieSocketAddress, ctx); |
| toSend.release(); |
| return; |
| } |
| boolean fenced = fencedLedgers.contains(ledgerId); |
| if (fenced && !isRecoveryAdd) { |
| callback.writeComplete(BKException.Code.LedgerFencedException, |
| ledgerId, entryId, bookieSocketAddress, ctx); |
| } else { |
| if (failedBookies.contains(bookieSocketAddress)) { |
| callback.writeComplete(NoBookieAvailableException, |
| ledgerId, entryId, bookieSocketAddress, ctx); |
| toSend.release(); |
| return; |
| } |
| if (getMockLedgerContentsInBookie(ledgerId, bookieSocketAddress).isEmpty()) { |
| registerMockEntryForRead(ledgerId, BookieProtocol.LAST_ADD_CONFIRMED, |
| bookieSocketAddress, new byte[0], BookieProtocol.INVALID_ENTRY_ID); |
| } |
| registerMockEntryForRead(ledgerId, entryId, bookieSocketAddress, entry, ledgerId); |
| callback.writeComplete(BKException.Code.OK, |
| ledgerId, entryId, bookieSocketAddress, ctx); |
| } |
| toSend.release(); |
| }); |
| |
| return null; |
| }); |
| |
| stub.when(bookieClient).addEntry(any(BookieSocketAddress.class), |
| anyLong(), any(byte[].class), |
| anyLong(), any(ByteBufList.class), |
| any(BookkeeperInternalCallbacks.WriteCallback.class), |
| any(), anyInt(), anyBoolean(), any(EnumSet.class)); |
| } |
| |
| @SuppressWarnings("unchecked") |
| protected void setupBookieClientForceLedger() { |
| final Stubber stub = doAnswer(invokation -> { |
| Object[] args = invokation.getArguments(); |
| BookieSocketAddress bookieSocketAddress = (BookieSocketAddress) args[0]; |
| long ledgerId = (Long) args[1]; |
| BookkeeperInternalCallbacks.ForceLedgerCallback callback = |
| (BookkeeperInternalCallbacks.ForceLedgerCallback) args[2]; |
| Object ctx = args[3]; |
| |
| Runnable activity = () -> { |
| executor.executeOrdered(ledgerId, () -> { |
| if (failedBookies.contains(bookieSocketAddress)) { |
| callback.forceLedgerComplete(NoBookieAvailableException, ledgerId, bookieSocketAddress, ctx); |
| return; |
| } |
| callback.forceLedgerComplete(BKException.Code.OK, ledgerId, bookieSocketAddress, ctx); |
| }); |
| }; |
| if (suspendedBookiesForForceLedgerAcks.contains(bookieSocketAddress)) { |
| List<Runnable> queue = deferredBookieForceLedgerResponses.computeIfAbsent(bookieSocketAddress, |
| (k) -> new CopyOnWriteArrayList<>()); |
| queue.add(activity); |
| } else { |
| activity.run(); |
| } |
| return null; |
| }); |
| |
| stub.when(bookieClient).forceLedger(any(BookieSocketAddress.class), |
| anyLong(), |
| any(BookkeeperInternalCallbacks.ForceLedgerCallback.class), |
| any()); |
| } |
| |
| } |
