Java tutorial
/* * Licensed 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.clients.impl.kv; import static org.apache.bookkeeper.stream.protocol.ProtocolConstants.DEFAULT_STREAM_CONF; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.same; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufUtil; import io.netty.buffer.Unpooled; import java.util.List; import java.util.NavigableMap; import java.util.Optional; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ThreadLocalRandom; import org.apache.bookkeeper.api.kv.PTable; import org.apache.bookkeeper.api.kv.Txn; import org.apache.bookkeeper.api.kv.impl.options.OptionFactoryImpl; import org.apache.bookkeeper.api.kv.options.DeleteOption; import org.apache.bookkeeper.api.kv.options.IncrementOption; import org.apache.bookkeeper.api.kv.options.OptionFactory; import org.apache.bookkeeper.api.kv.options.PutOption; import org.apache.bookkeeper.api.kv.options.RangeOption; import org.apache.bookkeeper.clients.exceptions.ClientException; import org.apache.bookkeeper.clients.impl.internal.api.HashStreamRanges; import org.apache.bookkeeper.clients.impl.internal.api.MetaRangeClient; import org.apache.bookkeeper.clients.impl.internal.api.StorageServerClientManager; import org.apache.bookkeeper.clients.impl.routing.RangeRouter; import org.apache.bookkeeper.clients.utils.ClientConstants; import org.apache.bookkeeper.common.concurrent.FutureUtils; import org.apache.bookkeeper.common.router.HashRouter; import org.apache.bookkeeper.common.util.Bytes; import org.apache.bookkeeper.common.util.OrderedScheduler; import org.apache.bookkeeper.stream.proto.RangeKeyType; import org.apache.bookkeeper.stream.proto.RangeProperties; import org.apache.bookkeeper.stream.proto.StreamProperties; import org.apache.bookkeeper.stream.protocol.util.ProtoUtils; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TestName; /** * Unit test for {@link PByteBufTableImpl}. */ public class TestPByteBufTableImpl { private static final long streamId = 12345L; @Rule public TestName runtime = new TestName(); private final HashStreamRanges streamRanges1 = prepareRanges(streamId, 4, 0); private final HashStreamRanges streamRanges2 = prepareRanges(streamId, 8, 4L); private final HashStreamRanges streamRanges3 = prepareRanges(streamId, 80, 12L); private final HashRouter<Integer> router = new HashRouter<Integer>() { private static final long serialVersionUID = -9119055960554608491L; private final List<Long> keys = Lists.newArrayList(streamRanges3.getRanges().keySet()); @Override public Long getRoutingKey(Integer key) { int idx; if (null == key) { idx = ThreadLocalRandom.current().nextInt(keys.size()); } else { idx = key % keys.size(); } return keys.get(idx); } }; private final StreamProperties streamProps = StreamProperties.newBuilder().setStorageContainerId(12345L) .setStreamConf(DEFAULT_STREAM_CONF).setStreamId(streamId).setStreamName("test-stream").build(); private final MetaRangeClient mockMetaRangeClient = mock(MetaRangeClient.class); private final StorageServerClientManager mockClientManager = mock(StorageServerClientManager.class); private final OptionFactory<ByteBuf> optionFactory = new OptionFactoryImpl<>(); private OrderedScheduler scheduler; private static HashStreamRanges prepareRanges(long streamId, int numRanges, long nextRangeId) { List<RangeProperties> ranges = ProtoUtils.split(streamId, numRanges, nextRangeId, (sid, rid) -> 1L); NavigableMap<Long, RangeProperties> rangeMap = Maps.newTreeMap(); for (RangeProperties props : ranges) { rangeMap.put(props.getStartHashKey(), props); } return HashStreamRanges.ofHash(RangeKeyType.HASH, rangeMap); } @Before public void setUp() { when(mockClientManager.openMetaRangeClient(any(StreamProperties.class))).thenReturn(mockMetaRangeClient); scheduler = OrderedScheduler.newSchedulerBuilder().numThreads(1).name("test-scheduler").build(); } @Test public void testInitializeFailureOnGetActiveRanges() { ClientException cause = new ClientException("test-cause"); when(mockMetaRangeClient.getActiveDataRanges()).thenReturn(FutureUtils.exception(cause)); PByteBufTableImpl table = new PByteBufTableImpl(runtime.getMethodName(), streamProps, mockClientManager, scheduler.chooseThread(1), ClientConstants.DEFAULT_INFINIT_BACKOFF_POLICY); try { FutureUtils.result(table.initialize()); fail("Should fail initializing the table with exception " + cause); } catch (Exception e) { assertEquals(cause, e); } } @SuppressWarnings("unchecked") @Test public void testBasicOperations() throws Exception { when(mockMetaRangeClient.getActiveDataRanges()).thenReturn(FutureUtils.value(streamRanges1)); ConcurrentMap<Long, PTable<ByteBuf, ByteBuf>> tableRanges = Maps.newConcurrentMap(); for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { tableRanges.put(rangeProps.getRangeId(), mock(PTable.class)); } RangeRouter<ByteBuf> mockRouter = mock(RangeRouter.class); when(mockRouter.getRange(any(ByteBuf.class))).thenAnswer(invocationOnMock -> { ByteBuf key = invocationOnMock.getArgument(0); byte[] keyData = ByteBufUtil.getBytes(key); return Bytes.toLong(keyData, 0); }); TableRangeFactory<ByteBuf, ByteBuf> trFactory = (streamProps1, rangeProps, executor, opFactory, resultFactory, kvFactory) -> tableRanges.get(rangeProps.getRangeId()); PByteBufTableImpl table = new PByteBufTableImpl(runtime.getMethodName(), streamProps, mockClientManager, scheduler.chooseThread(), trFactory, Optional.of(mockRouter)); assertEquals(0, table.getTableRanges().size()); verify(mockRouter, times(0)).setRanges(any(HashStreamRanges.class)); // initialize the table assertTrue(table == FutureUtils.result(table.initialize())); verify(mockRouter, times(1)).setRanges(eq(streamRanges1)); assertEquals(4, table.getTableRanges().size()); // test get for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { ByteBuf pkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); ByteBuf lkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); try (RangeOption<ByteBuf> option = optionFactory.newRangeOption().build()) { table.get(pkey, lkey, option); verify(tableRanges.get(rangeProps.getRangeId()), times(1)).get(eq(pkey), eq(lkey), eq(option)); } } // test put for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { ByteBuf pkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); ByteBuf lkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); ByteBuf value = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); try (PutOption<ByteBuf> option = optionFactory.newPutOption().build()) { table.put(pkey, lkey, value, option); verify(tableRanges.get(rangeProps.getRangeId()), times(1)).put(eq(pkey), eq(lkey), eq(value), eq(option)); } } // test increment for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { ByteBuf pkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); ByteBuf lkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); long amount = 100L; try (IncrementOption<ByteBuf> option = optionFactory.newIncrementOption().build()) { table.increment(pkey, lkey, amount, option); verify(tableRanges.get(rangeProps.getRangeId()), times(1)).increment(eq(pkey), eq(lkey), eq(amount), same(option)); } } // test delete for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { ByteBuf pkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); ByteBuf lkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); try (DeleteOption<ByteBuf> option = optionFactory.newDeleteOption().build()) { table.delete(pkey, lkey, option); verify(tableRanges.get(rangeProps.getRangeId()), times(1)).delete(eq(pkey), eq(lkey), eq(option)); } } // test txn for (RangeProperties rangeProps : streamRanges1.getRanges().values()) { ByteBuf pkey = Unpooled.wrappedBuffer(Bytes.toBytes(rangeProps.getRangeId())); Txn<ByteBuf, ByteBuf> txn = table.txn(pkey); verify(tableRanges.get(rangeProps.getRangeId()), times(1)).txn(eq(pkey)); } } }