tests/integration/cluster/src/test/java/org/apache/bookkeeper/tests/integration/stream/TableClientSimpleTest.java

/*
 * 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.tests.integration.stream;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.bookkeeper.common.concurrent.FutureUtils.result;
import static org.apache.bookkeeper.stream.protocol.ProtocolConstants.DEFAULT_STREAM_CONF;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.Unpooled;
import java.util.List;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.api.StorageClient;
import org.apache.bookkeeper.api.kv.PTable;
import org.apache.bookkeeper.api.kv.exceptions.KvApiException;
import org.apache.bookkeeper.api.kv.result.Code;
import org.apache.bookkeeper.api.kv.result.KeyValue;
import org.apache.bookkeeper.clients.admin.StorageAdminClient;
import org.apache.bookkeeper.clients.config.StorageClientSettings;
import org.apache.bookkeeper.common.testing.annotations.FlakyTest;
import org.apache.bookkeeper.stream.proto.NamespaceConfiguration;
import org.apache.bookkeeper.stream.proto.NamespaceProperties;
import org.apache.bookkeeper.stream.proto.StorageType;
import org.apache.bookkeeper.stream.proto.StreamConfiguration;
import org.apache.bookkeeper.stream.proto.StreamProperties;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.rules.TestName;

/**
 * Integration test for table service.
 */
@Slf4j
public class TableClientSimpleTest extends StreamClusterTestBase {

    @Rule
    public final TestName testName = new TestName();

    private final String namespace = "test_namespace";
    private StorageAdminClient adminClient;
    private StorageClient storageClient;

    @Before
    public void setup() {
        StorageClientSettings settings = newStorageClientSettings();
        String namespace = "test_namespace";
        adminClient = createStorageAdminClient(settings);
        storageClient = createStorageClient(settings, namespace);
    }

    @After
    public void teardown() {
        if (null != adminClient) {
            adminClient.close();
        }
        if (null != storageClient) {
            storageClient.close();
        }
    }

    private static ByteBuf getLKey(int i) {
        return Unpooled.wrappedBuffer(String.format("test-lkey-%06d", i).getBytes(UTF_8));
    }

    private static ByteBuf getValue(int i) {
        return Unpooled.wrappedBuffer(String.format("test-val-%06d", i).getBytes(UTF_8));
    }

    @FlakyTest("https://github.com/apache/bookkeeper/issues/1440")
    public void testTableSimpleAPI() throws Exception {
        // Create a namespace
        NamespaceConfiguration nsConf = NamespaceConfiguration.newBuilder()
            .setDefaultStreamConf(DEFAULT_STREAM_CONF)
            .build();
        NamespaceProperties nsProps = result(adminClient.createNamespace(namespace, nsConf));
        assertEquals(namespace, nsProps.getNamespaceName());
        assertEquals(nsConf.getDefaultStreamConf(), nsProps.getDefaultStreamConf());

        // Create a stream
        String streamName = testName.getMethodName() + "_stream";
        StreamConfiguration streamConf = StreamConfiguration.newBuilder(DEFAULT_STREAM_CONF)
            .setStorageType(StorageType.TABLE)
            .build();
        StreamProperties streamProps = result(
            adminClient.createStream(namespace, streamName, streamConf));
        assertEquals(streamName, streamProps.getStreamName());

        // Open the table
        PTable<ByteBuf, ByteBuf> table = result(storageClient.openPTable(streamName));
        byte[] rKey = "routing-key".getBytes(UTF_8);
        byte[] lKey = "testing-key".getBytes(UTF_8);
        byte[] value1 = "testing-value-1".getBytes(UTF_8);
        byte[] value2 = "testing-value-2".getBytes(UTF_8);

        // put first key
        ByteBuf rKeyBuf = Unpooled.wrappedBuffer(rKey);
        ByteBuf lKeyBuf = Unpooled.wrappedBuffer(lKey);
        ByteBuf valBuf1 = Unpooled.wrappedBuffer(value1);
        ByteBuf valBuf2 = Unpooled.wrappedBuffer(value2);

        // normal put
        assertNull(result(
            table.put(rKeyBuf, lKeyBuf, valBuf1)));

        // putIfAbsent failure
        assertArrayEquals(
            value1,
            ByteBufUtil.getBytes(result(table.putIfAbsent(rKeyBuf, lKeyBuf, valBuf2))));

        // delete failure
        assertFalse(
            result(table.delete(rKeyBuf, lKeyBuf, valBuf2)));

        // delete success
        assertTrue(
            result(table.delete(rKeyBuf, lKeyBuf, valBuf1)));

        // get
        assertNull(
            result(table.get(rKeyBuf, lKeyBuf)));

        // putIfAbsent success
        assertNull(
            result(table.putIfAbsent(rKeyBuf, lKeyBuf, valBuf2)));

        // get returns val2
        assertArrayEquals(
            value2,
            ByteBufUtil.getBytes(result(table.get(rKeyBuf, lKeyBuf))));

        // vPut failure
        try {
            result(
                table.vPut(rKeyBuf, lKeyBuf, valBuf1, 9999L));
            fail("Should fail vPut if the version doesn't match");
        } catch (KvApiException e) {
            assertEquals(Code.BAD_REVISION, e.getCode());
        }

        // vPut success
        assertEquals(
            1L,
            result(
                table.vPut(rKeyBuf, lKeyBuf, valBuf1, 0L)).longValue());

        // vDelete failure
        try {
            result(
                table.vDelete(rKeyBuf, lKeyBuf, 9999L));
            fail("Should fail vDelete if the version doesn't match");
        } catch (KvApiException e) {
            assertEquals(Code.BAD_REVISION, e.getCode());
        }

        // vDelete success
        try (KeyValue<ByteBuf, ByteBuf> prevKv = result(
            table.vDelete(rKeyBuf, lKeyBuf, 1L))) {
            assertNotNull(prevKv);
            assertEquals(1L, prevKv.version());
            assertArrayEquals(
                value1,
                ByteBufUtil.getBytes(prevKv.value()));
        }

        // write a range of key
        int numKvs = 100;
        rKeyBuf = Unpooled.wrappedBuffer("test-key".getBytes(UTF_8));
        for (int i = 0; i < numKvs; i++) {
            lKeyBuf = getLKey(i);
            valBuf1 = getValue(i);
            result(table.put(rKeyBuf, lKeyBuf, valBuf1));
        }

        // get ranges
        ByteBuf lStartKey = getLKey(20);
        ByteBuf lEndKey = getLKey(50);
        List<KeyValue<ByteBuf, ByteBuf>> kvs = result(
            table.range(rKeyBuf, lStartKey, lEndKey));
        assertEquals(31, kvs.size());
        int i = 20;
        for (KeyValue<ByteBuf, ByteBuf> kvPair : kvs) {
            assertEquals(getLKey(i), kvPair.key());
            assertEquals(getValue(i), kvPair.value());
            ++i;
            kvPair.close();
        }
        assertEquals(51, i);

        // delete range
        kvs = result(
            table.deleteRange(rKeyBuf, lStartKey, lEndKey));
        assertEquals(31, kvs.size());
        i = 20;
        for (KeyValue<ByteBuf, ByteBuf> kvPair : kvs) {
            assertEquals(getLKey(i), kvPair.key());
            assertEquals(getValue(i), kvPair.value());
            ++i;
            kvPair.close();
        }
        assertEquals(51, i);

        // get ranges again
        kvs = result(table.range(rKeyBuf, lStartKey, lEndKey));
        assertTrue(kvs.isEmpty());

        byte[] lIncrKey = "test-incr-lkey".getBytes(UTF_8);
        ByteBuf lIncrKeyBuf = Unpooled.wrappedBuffer(lIncrKey);

        // test increment
        for (int j = 0; j < 5; j++) {
            result(table.increment(rKeyBuf, lIncrKeyBuf, 100L));
            long number = result(table.getNumber(rKeyBuf, lIncrKeyBuf));
            assertEquals(100L * (j + 1), number);
        }

        for (int j = 5; j < 10; j++) {
            long number = result(table.incrementAndGet(rKeyBuf, lIncrKeyBuf, 100L));
            assertEquals(100L * (j + 1), number);
        }
    }

    public void testTableTtl() throws Exception {
        final String ns = namespace + "_ttl";
        final int ttlSeconds = 5;

        result(adminClient.createNamespace(ns, NamespaceConfiguration.newBuilder()
                .setDefaultStreamConf(DEFAULT_STREAM_CONF)
                .build()));

        final String streamName = testName.getMethodName() + "_stream";
        result(adminClient.createStream(ns, streamName, StreamConfiguration.newBuilder(DEFAULT_STREAM_CONF)
                        .setStorageType(StorageType.TABLE)
                        .setTtlSeconds(ttlSeconds)
                        .build()));

        final PTable<ByteBuf, ByteBuf> table = result(storageClient.openPTable(streamName));

        // put first key
        final byte[] rKey = "rKey".getBytes(UTF_8);
        final byte[] lKey = "lKey".getBytes(UTF_8);
        final byte[] val = "val".getBytes(UTF_8);

        final ByteBuf rKeyBuf = Unpooled.wrappedBuffer(rKey);
        final ByteBuf lKeyBuf = Unpooled.wrappedBuffer(lKey);
        final ByteBuf valBuf = Unpooled.wrappedBuffer(val);

        result(table.put(rKeyBuf, lKeyBuf, valBuf));

        // Ensure we can read the value back.  this may be flaky if the test gets suspended for too long.
        assertArrayEquals(val, ByteBufUtil.getBytes(result(table.get(rKeyBuf, lKeyBuf))));

        // Wait for the value to expire, and ensure we can't read it.
        Thread.sleep(ttlSeconds + 1);
        assertNull(ByteBufUtil.getBytes(result(table.get(rKeyBuf, lKeyBuf))));
    }
}