/**
 * 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 com.twitter.distributedlog;

import java.net.URI;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import com.twitter.distributedlog.exceptions.AlreadyTruncatedTransactionException;
import com.twitter.distributedlog.exceptions.BKTransmitException;
import com.twitter.distributedlog.exceptions.LogEmptyException;
import com.twitter.distributedlog.exceptions.LogNotFoundException;
import com.twitter.distributedlog.exceptions.LogReadException;
import com.twitter.distributedlog.impl.ZKLogSegmentMetadataStore;
import com.twitter.distributedlog.logsegment.LogSegmentMetadataStore;
import com.twitter.distributedlog.util.FutureUtils;
import com.twitter.distributedlog.util.OrderedScheduler;
import com.twitter.distributedlog.util.Utils;
import org.apache.bookkeeper.client.BKException;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.twitter.distributedlog.callback.LogSegmentListener;
import com.twitter.distributedlog.exceptions.EndOfStreamException;
import com.twitter.distributedlog.exceptions.InvalidStreamNameException;
import com.twitter.distributedlog.exceptions.LogRecordTooLongException;
import com.twitter.distributedlog.exceptions.OwnershipAcquireFailedException;
import com.twitter.distributedlog.exceptions.TransactionIdOutOfOrderException;
import com.twitter.distributedlog.impl.metadata.ZKLogMetadata;
import com.twitter.distributedlog.metadata.BKDLConfig;
import com.twitter.distributedlog.metadata.MetadataUpdater;
import com.twitter.distributedlog.metadata.LogSegmentMetadataStoreUpdater;
import com.twitter.distributedlog.namespace.DistributedLogNamespace;
import com.twitter.distributedlog.namespace.DistributedLogNamespaceBuilder;
import com.twitter.distributedlog.subscription.SubscriptionStateStore;
import com.twitter.distributedlog.subscription.SubscriptionsStore;
import com.twitter.util.Await;
import com.twitter.util.Duration;
import com.twitter.util.Future;

import static org.junit.Assert.*;
import static org.junit.Assert.assertEquals;

public class TestBKDistributedLogManager extends TestDistributedLogBase {
    static final Logger LOG = LoggerFactory.getLogger(TestBKDistributedLogManager.class);

    @Rule
    public TestName testNames = new TestName();

    private static final long DEFAULT_SEGMENT_SIZE = 1000;

    private void testNonPartitionedWritesInternal(String name, DistributedLogConfiguration conf) throws Exception {
        BKDistributedLogManager dlm = createNewDLM(conf, name);

        long txid = 1;
        for (long i = 0; i < 3; i++) {
            long start = txid;
            BKSyncLogWriter writer = dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                writer.write(DLMTestUtil.getLogRecordInstance(txid++));
            }
            BKLogSegmentWriter perStreamLogWriter = writer.getCachedLogWriter();
            writer.closeAndComplete();
            BKLogWriteHandler blplm = dlm.createWriteHandler(true);
            assertNotNull(zkc.exists(blplm.completedLedgerZNode(start, txid - 1,
                                                                perStreamLogWriter.getLogSegmentSequenceNumber()), false));
            FutureUtils.result(blplm.asyncClose());
        }

        LogWriter writer = dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            writer.write(DLMTestUtil.getLogRecordInstance(txid++));
        }
        writer.setReadyToFlush();
        writer.flushAndSync();
        writer.close();

        LogReader reader = dlm.getInputStream(1);
        long numTrans = 0;
        LogRecord record = reader.readNext(false);
        long lastTxId = -1;
        while (null != record) {
            DLMTestUtil.verifyLogRecord(record);
            assert (lastTxId < record.getTransactionId());
            lastTxId = record.getTransactionId();
            numTrans++;
            record = reader.readNext(false);
        }
        reader.close();
        assertEquals((txid - 1), numTrans);
    }

    @Test(timeout = 60000)
    public void testSimpleWrite() throws Exception {
        BKDistributedLogManager dlm = createNewDLM(conf, "distrlog-simplewrite");
        BKSyncLogWriter out = dlm.startLogSegmentNonPartitioned();
        for (long i = 1; i <= 100; i++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(i);
            out.write(op);
        }
        BKLogSegmentWriter perStreamLogWriter = out.getCachedLogWriter();
        out.closeAndComplete();

        BKLogWriteHandler blplm = dlm.createWriteHandler(true);
        assertNotNull(zkc.exists(blplm.completedLedgerZNode(1, 100,
                perStreamLogWriter.getLogSegmentSequenceNumber()), false));
        FutureUtils.result(blplm.asyncClose());
    }

    @Test(timeout = 60000)
    public void testNumberOfTransactions() throws Exception {
        String name = "distrlog-txncount";
        DistributedLogManager dlm = createNewDLM(conf, name);
        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long i = 1; i <= 100; i++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(i);
            out.write(op);
        }
        out.closeAndComplete();

        long numTrans = DLMTestUtil.getNumberofLogRecords(createNewDLM(conf, name), 1);
        assertEquals(100, numTrans);
        dlm.close();
    }

    @Test(timeout = 60000)
    public void testSanityCheckTxnID() throws Exception {
        String name = "distrlog-sanity-check-txnid";
        BKDistributedLogManager dlm = createNewDLM(conf, name);
        BKSyncLogWriter out = dlm.startLogSegmentNonPartitioned();
        long txid = 1;
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.closeAndComplete();

        BKSyncLogWriter out1 = dlm.startLogSegmentNonPartitioned();
        LogRecord op1 = DLMTestUtil.getLogRecordInstance(1);
        try {
            out1.write(op1);
            fail("Should fail writing lower txn id if sanityCheckTxnID is enabled.");
        } catch (TransactionIdOutOfOrderException tioooe) {
            // expected
        }
        out1.closeAndComplete();
        dlm.close();

        DLMTestUtil.updateBKDLConfig(bkutil.getUri(), bkutil.getZkServers(), bkutil.getBkLedgerPath(), false);
        LOG.info("Disable sanity check txn id.");
        BKDLConfig.clearCachedDLConfigs();

        DistributedLogConfiguration newConf = new DistributedLogConfiguration();
        newConf.addConfiguration(conf);
        BKDistributedLogManager newDLM = createNewDLM(newConf, name);
        BKSyncLogWriter out2 = newDLM.startLogSegmentNonPartitioned();
        LogRecord op2 = DLMTestUtil.getLogRecordInstance(1);
        out2.write(op2);
        out2.closeAndComplete();
        newDLM.close();

        DLMTestUtil.updateBKDLConfig(bkutil.getUri(), bkutil.getZkServers(), bkutil.getBkLedgerPath(), true);
        LOG.info("Enable sanity check txn id.");
        BKDLConfig.clearCachedDLConfigs();

        DistributedLogConfiguration conf3 = new DistributedLogConfiguration();
        conf3.addConfiguration(conf);
        BKDistributedLogManager dlm3 = createNewDLM(newConf, name);
        BKSyncLogWriter out3 = dlm3.startLogSegmentNonPartitioned();
        LogRecord op3 = DLMTestUtil.getLogRecordInstance(1);
        try {
            out3.write(op3);
            fail("Should fail writing lower txn id if sanityCheckTxnID is enabled.");
        } catch (TransactionIdOutOfOrderException tioooe) {
            // expected
        }
        out3.closeAndComplete();
        dlm3.close();
    }

    @Test(timeout = 60000)
    public void testContinuousReaders() throws Exception {
        String name = "distrlog-continuous";
        BKDistributedLogManager dlm = createNewDLM(conf, name);
        long txid = 1;
        for (long i = 0; i < 3; i++) {
            long start = txid;
            BKSyncLogWriter out = dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
                out.write(op);
            }
            BKLogSegmentWriter perStreamLogWriter = out.getCachedLogWriter();
            out.closeAndComplete();
            BKLogWriteHandler blplm = dlm.createWriteHandler(true);

            assertNotNull(
                zkc.exists(blplm.completedLedgerZNode(start, txid - 1,
                                                      perStreamLogWriter.getLogSegmentSequenceNumber()), false));
            FutureUtils.result(blplm.asyncClose());
        }

        BKSyncLogWriter out = dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.setReadyToFlush();
        out.flushAndSync();
        out.close();
        dlm.close();

        dlm = createNewDLM(conf, name);

        LogReader reader = dlm.getInputStream(1);
        long numTrans = 0;
        LogRecord record = reader.readNext(false);
        while (null != record) {
            DLMTestUtil.verifyLogRecord(record);
            numTrans++;
            record = reader.readNext(false);
        }
        assertEquals((txid - 1), numTrans);
        assertEquals(txid - 1, dlm.getLogRecordCount());
        reader.close();
        dlm.close();
    }

    /**
     * Create a bkdlm namespace, write a journal from txid 1, close stream.
     * Try to create a new journal from txid 1. Should throw an exception.
     */
    @Test(timeout = 60000)
    public void testWriteRestartFrom1() throws Exception {
        DistributedLogManager dlm = createNewDLM(conf, "distrlog-restartFrom1");
        long txid = 1;
        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.closeAndComplete();

        txid = 1;
        try {
            out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            out.write(DLMTestUtil.getLogRecordInstance(txid));
            fail("Shouldn't be able to start another journal from " + txid
                + " when one already exists");
        } catch (Exception ioe) {
            LOG.info("Caught exception as expected", ioe);
        } finally {
            out.close();
        }

        // test border case
        txid = DEFAULT_SEGMENT_SIZE - 1;
        try {
            out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            out.write(DLMTestUtil.getLogRecordInstance(txid));
            fail("Shouldn't be able to start another journal from " + txid
                + " when one already exists");
        } catch (TransactionIdOutOfOrderException rste) {
            LOG.info("Caught exception as expected", rste);
        } finally {
            out.close();
        }

        // open journal continuing from before
        txid = DEFAULT_SEGMENT_SIZE + 1;
        out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        assertNotNull(out);

        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.closeAndComplete();

        // open journal arbitarily far in the future
        txid = DEFAULT_SEGMENT_SIZE * 4;
        out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        out.write(DLMTestUtil.getLogRecordInstance(txid));
        out.close();
        dlm.close();
    }

    @Test(timeout = 60000)
    public void testTwoWriters() throws Exception {
        DLMTestUtil.BKLogPartitionWriteHandlerAndClients bkdlm1 =
                createNewBKDLM(conf, "distrlog-dualWriter");
        try {
             createNewBKDLM(conf, "distrlog-dualWriter");
            fail("Shouldn't have been able to open the second writer");
        } catch (OwnershipAcquireFailedException ioe) {
            assertEquals(ioe.getCurrentOwner(), DistributedLogConstants.UNKNOWN_CLIENT_ID);
        }

        bkdlm1.close();
    }

    @Test(timeout = 60000)
    public void testTwoWritersOnLockDisabled() throws Exception {
        DistributedLogConfiguration confLocal = new DistributedLogConfiguration();
        confLocal.addConfiguration(conf);
        confLocal.setOutputBufferSize(0);
        confLocal.setWriteLockEnabled(false);
        String name = "distrlog-two-writers-lock-disabled";
        DistributedLogManager manager = createNewDLM(confLocal, name);
        AsyncLogWriter writer1 = FutureUtils.result(manager.openAsyncLogWriter());
        FutureUtils.result(writer1.write(DLMTestUtil.getLogRecordInstance(1L)));
        AsyncLogWriter writer2 = FutureUtils.result(manager.openAsyncLogWriter());
        FutureUtils.result(writer2.write(DLMTestUtil.getLogRecordInstance(2L)));

        // write a record to writer 1 again
        try {
            FutureUtils.result(writer1.write(DLMTestUtil.getLogRecordInstance(3L)));
            fail("Should fail writing record to writer 1 again as writer 2 took over the ownership");
        } catch (BKTransmitException bkte) {
            assertEquals(BKException.Code.LedgerFencedException, bkte.getBKResultCode());
        }
    }

    @Test(timeout = 60000)
    public void testSimpleRead() throws Exception {
        String name = "distrlog-simpleread";
        DistributedLogManager dlm = createNewDLM(conf, name);
        final long numTransactions = 10000;
        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long i = 1; i <= numTransactions; i++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(i);
            out.write(op);
        }
        out.closeAndComplete();

        assertEquals(numTransactions, DLMTestUtil.getNumberofLogRecords(createNewDLM(conf, name), 1));
        dlm.close();
    }

    @Test(timeout = 60000)
    public void testNumberOfTransactionsWithInprogressAtEnd() throws Exception {
        String name = "distrlog-inprogressAtEnd";
        DistributedLogManager dlm = createNewDLM(conf, name);
        long txid = 1;
        for (long i = 0; i < 3; i++) {
            long start = txid;
            BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
                out.write(op);
            }
            BKLogSegmentWriter perStreamLogWriter = out.getCachedLogWriter();
            out.closeAndComplete();
            BKLogWriteHandler blplm = ((BKDistributedLogManager) (dlm)).createWriteHandler(true);
            assertNotNull(
                zkc.exists(blplm.completedLedgerZNode(start, txid - 1,
                                                      perStreamLogWriter.getLogSegmentSequenceNumber()), false));
            FutureUtils.result(blplm.asyncClose());
        }
        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.setReadyToFlush();
        out.flushAndSync();
        out.close();

        long numTrans = DLMTestUtil.getNumberofLogRecords(createNewDLM(conf, name), 1);
        assertEquals((txid - 1), numTrans);
    }

    @Test(timeout = 60000)
    public void testContinuousReaderBulk() throws Exception {
        String name = "distrlog-continuous-bulk";
        DistributedLogManager dlm = createNewDLM(conf, name);
        long txid = 1;
        for (long i = 0; i < 3; i++) {
            BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
                out.write(op);
            }
            out.closeAndComplete();
        }

        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.setReadyToFlush();
        out.flushAndSync();
        out.close();
        dlm.close();

        dlm = createNewDLM(conf, name);

        LogReader reader = dlm.getInputStream(1);
        long numTrans = 0;
        List<LogRecordWithDLSN> recordList = reader.readBulk(false, 13);
        long lastTxId = -1;
        while (!recordList.isEmpty()) {
            for (LogRecord record : recordList) {
                assert (lastTxId < record.getTransactionId());
                lastTxId = record.getTransactionId();
                DLMTestUtil.verifyLogRecord(record);
                numTrans++;
            }
            recordList = reader.readBulk(false, 13);
        }
        reader.close();
        assertEquals((txid - 1), numTrans);
    }

    @Test(timeout = 60000)
    public void testContinuousReadersWithEmptyLedgers() throws Exception {
        String name = "distrlog-continuous-emptyledgers";
        DistributedLogManager dlm = createNewDLM(conf, name);
        long txid = 1;
        for (long i = 0; i < 3; i++) {
            long start = txid;
            BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
                out.write(op);
            }
            BKLogSegmentWriter writer = out.getCachedLogWriter();
            out.closeAndComplete();
            BKLogWriteHandler blplm = ((BKDistributedLogManager) (dlm)).createWriteHandler(true);

            assertNotNull(
                zkc.exists(blplm.completedLedgerZNode(start, txid - 1,
                                                      writer.getLogSegmentSequenceNumber()), false));
            BKLogSegmentWriter perStreamLogWriter = blplm.startLogSegment(txid - 1);
            blplm.completeAndCloseLogSegment(perStreamLogWriter.getLogSegmentSequenceNumber(),
                    perStreamLogWriter.getLogSegmentId(), txid - 1, txid - 1, 0);
            assertNotNull(
                zkc.exists(blplm.completedLedgerZNode(txid - 1, txid - 1,
                                                      perStreamLogWriter.getLogSegmentSequenceNumber()), false));
            FutureUtils.result(blplm.asyncClose());
        }

        BKSyncLogWriter out = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            LogRecord op = DLMTestUtil.getLogRecordInstance(txid++);
            out.write(op);
        }
        out.setReadyToFlush();
        out.flushAndSync();
        out.close();
        dlm.close();

        dlm = createNewDLM(conf, name);

        AsyncLogReader asyncreader = dlm.getAsyncLogReader(DLSN.InvalidDLSN);
        long numTrans = 0;
        LogRecordWithDLSN record = Await.result(asyncreader.readNext());
        while (null != record) {
            DLMTestUtil.verifyLogRecord(record);
            numTrans++;
            if (numTrans >= (txid - 1)) {
                break;
            }
            record = Await.result(asyncreader.readNext());
        }
        assertEquals((txid - 1), numTrans);
        Utils.close(asyncreader);

        LogReader reader = dlm.getInputStream(1);
        numTrans = 0;
        record = reader.readNext(false);
        while (null != record) {
            DLMTestUtil.verifyLogRecord(record);
            numTrans++;
            record = reader.readNext(false);
        }
        assertEquals((txid - 1), numTrans);
        reader.close();
        assertEquals(txid - 1, dlm.getLogRecordCount());
        dlm.close();
    }

    @Test(timeout = 60000)
    public void testNonPartitionedWrites() throws Exception {
        String name = "distrlog-non-partitioned-bulk";
        testNonPartitionedWritesInternal(name, conf);
    }

    @Test(timeout = 60000)
    public void testCheckLogExists() throws Exception {
        String name = "distrlog-check-log-exists";
        DistributedLogManager dlm = createNewDLM(conf, name);

        long txid = 1;
        LogWriter writer = dlm.startLogSegmentNonPartitioned();
        for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
            writer.write(DLMTestUtil.getLogRecordInstance(txid++));
        }
        writer.setReadyToFlush();
        writer.flushAndSync();
        writer.close();
        dlm.createOrUpdateMetadata(name.getBytes());
        assertEquals(name, new String(dlm.getMetadata()));

        URI uri = createDLMURI("/" + name);
        BKDistributedLogNamespace namespace = BKDistributedLogNamespace.newBuilder()
                .conf(conf).uri(uri).build();
        assertTrue(namespace.logExists(name));
        assertFalse(namespace.logExists("non-existent-log"));
        URI nonExistentUri = createDLMURI("/" + "non-existent-ns");
        DistributedLogNamespace nonExistentNS = DistributedLogNamespaceBuilder.newBuilder()
                .conf(conf).uri(nonExistentUri).build();
        assertFalse(nonExistentNS.logExists(name));

        int logCount = 0;
        Iterator<String> logIter = namespace.getLogs();
        while(logIter.hasNext()) {
            String log = logIter.next();
            logCount++;
            assertEquals(name, log);
        }
        assertEquals(1, logCount);

        for(Map.Entry<String, byte[]> logEntry: namespace.enumerateLogsWithMetadataInNamespace().entrySet()) {
            assertEquals(name, new String(logEntry.getValue()));
        }
    }

    @Test(timeout = 60000)
    public void testMetadataAccessor() throws Exception {
        String name = "distrlog-metadata-accessor";
        MetadataAccessor metadata = DLMTestUtil.createNewMetadataAccessor(conf, name, createDLMURI("/" + name));
        assertEquals(name, metadata.getStreamName());
        metadata.createOrUpdateMetadata(name.getBytes());
        assertEquals(name, new String(metadata.getMetadata()));
        metadata.deleteMetadata();
        assertEquals(null, metadata.getMetadata());
    }

    @Test(timeout = 60000)
    @Deprecated
    public void testSubscriptionStateStore() throws Exception {
        String name = "distrlog-subscription-state";
        String subscriberId = "defaultSubscriber";
        DLSN commitPosition0 = new DLSN(4, 33, 5);
        DLSN commitPosition1 = new DLSN(4, 34, 5);
        DLSN commitPosition2 = new DLSN(5, 34, 5);

        DistributedLogManager dlm = createNewDLM(conf, name);
        SubscriptionStateStore store = dlm.getSubscriptionStateStore(subscriberId);
        assertEquals(Await.result(store.getLastCommitPosition()), DLSN.NonInclusiveLowerBound);
        Await.result(store.advanceCommitPosition(commitPosition1));
        assertEquals(Await.result(store.getLastCommitPosition()), commitPosition1);
        Await.result(store.advanceCommitPosition(commitPosition0));
        assertEquals(Await.result(store.getLastCommitPosition()), commitPosition1);
        Await.result(store.advanceCommitPosition(commitPosition2));
        assertEquals(Await.result(store.getLastCommitPosition()), commitPosition2);
        SubscriptionStateStore store1 = dlm.getSubscriptionStateStore(subscriberId);
        assertEquals(Await.result(store1.getLastCommitPosition()), commitPosition2);
    }

    @Test(timeout = 60000)
    public void testSubscriptionsStore() throws Exception {
        String name = "distrlog-subscriptions-store";
        String subscriber0 = "subscriber-0";
        String subscriber1 = "subscriber-1";
        String subscriber2 = "subscriber-2";

        DLSN commitPosition0 = new DLSN(4, 33, 5);
        DLSN commitPosition1 = new DLSN(4, 34, 5);
        DLSN commitPosition2 = new DLSN(5, 34, 5);
        DLSN commitPosition3 = new DLSN(6, 35, 6);

        DistributedLogManager dlm = createNewDLM(conf, name);

        SubscriptionsStore store = dlm.getSubscriptionsStore();

        // no data
        assertEquals(Await.result(store.getLastCommitPosition(subscriber0)), DLSN.NonInclusiveLowerBound);
        assertEquals(Await.result(store.getLastCommitPosition(subscriber1)), DLSN.NonInclusiveLowerBound);
        assertEquals(Await.result(store.getLastCommitPosition(subscriber2)), DLSN.NonInclusiveLowerBound);
        // empty
        assertTrue(Await.result(store.getLastCommitPositions()).isEmpty());

        // subscriber 0 advance
        Await.result(store.advanceCommitPosition(subscriber0, commitPosition0));
        assertEquals(commitPosition0, Await.result(store.getLastCommitPosition(subscriber0)));
        Map<String, DLSN> committedPositions = Await.result(store.getLastCommitPositions());
        assertEquals(1, committedPositions.size());
        assertEquals(commitPosition0, committedPositions.get(subscriber0));

        // subscriber 1 advance
        Await.result(store.advanceCommitPosition(subscriber1, commitPosition1));
        assertEquals(commitPosition1, Await.result(store.getLastCommitPosition(subscriber1)));
        committedPositions = Await.result(store.getLastCommitPositions());
        assertEquals(2, committedPositions.size());
        assertEquals(commitPosition0, committedPositions.get(subscriber0));
        assertEquals(commitPosition1, committedPositions.get(subscriber1));

        // subscriber 2 advance
        Await.result(store.advanceCommitPosition(subscriber2, commitPosition2));
        assertEquals(commitPosition2, Await.result(store.getLastCommitPosition(subscriber2)));
        committedPositions = Await.result(store.getLastCommitPositions());
        assertEquals(3, committedPositions.size());
        assertEquals(commitPosition0, committedPositions.get(subscriber0));
        assertEquals(commitPosition1, committedPositions.get(subscriber1));
        assertEquals(commitPosition2, committedPositions.get(subscriber2));

        // subscriber 2 advance again
        DistributedLogManager newDLM = createNewDLM(conf, name);
        SubscriptionsStore newStore = newDLM.getSubscriptionsStore();
        Await.result(newStore.advanceCommitPosition(subscriber2, commitPosition3));
        newStore.close();
        newDLM.close();

        committedPositions = Await.result(store.getLastCommitPositions());
        assertEquals(3, committedPositions.size());
        assertEquals(commitPosition0, committedPositions.get(subscriber0));
        assertEquals(commitPosition1, committedPositions.get(subscriber1));
        assertEquals(commitPosition3, committedPositions.get(subscriber2));

        dlm.close();

    }

    private long writeAndMarkEndOfStream(DistributedLogManager dlm, long txid) throws Exception {
        for (long i = 0; i < 3; i++) {
            long start = txid;
            BKSyncLogWriter writer = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE; j++) {
                writer.write(DLMTestUtil.getLogRecordInstance(txid++));
            }

            BKLogSegmentWriter perStreamLogWriter = writer.getCachedLogWriter();

            if (i < 2) {
                writer.closeAndComplete();
                BKLogWriteHandler blplm = ((BKDistributedLogManager) (dlm)).createWriteHandler(true);
                assertNotNull(zkc.exists(blplm.completedLedgerZNode(start, txid - 1,
                                                                    perStreamLogWriter.getLogSegmentSequenceNumber()), false));
                FutureUtils.result(blplm.asyncClose());
            } else {
                writer.markEndOfStream();
                BKLogWriteHandler blplm = ((BKDistributedLogManager) (dlm)).createWriteHandler(true);
                assertNotNull(zkc.exists(blplm.completedLedgerZNode(start, DistributedLogConstants.MAX_TXID,
                                                                    perStreamLogWriter.getLogSegmentSequenceNumber()), false));
                FutureUtils.result(blplm.asyncClose());
            }
        }
        return txid;
    }

    @Test(timeout = 60000)
    public void testMarkEndOfStream() throws Exception {
        String name = "distrlog-mark-end-of-stream";
        DistributedLogManager dlm = createNewDLM(conf, name);

        long txid = 1;
        txid = writeAndMarkEndOfStream(dlm, txid);

        LogReader reader = dlm.getInputStream(1);
        long numTrans = 0;
        boolean exceptionEncountered = false;
        LogRecord record = null;
        try {
            record = reader.readNext(false);
            long expectedTxId = 1;
            while (null != record) {
                DLMTestUtil.verifyLogRecord(record);
                assertEquals(expectedTxId, record.getTransactionId());
                expectedTxId++;
                numTrans++;
                record = reader.readNext(false);
            }
        } catch (EndOfStreamException exc) {
            LOG.info("Encountered EndOfStream on reading records after {}", record);
            exceptionEncountered = true;
        }
        assertEquals((txid - 1), numTrans);
        assertTrue(exceptionEncountered);
        exceptionEncountered = false;
        try {
            reader.readNext(false);
        } catch (EndOfStreamException exc) {
            exceptionEncountered = true;
        }
        assertTrue(exceptionEncountered);
        reader.close();
    }

    @Test(timeout = 60000)
    public void testWriteFailsAfterMarkEndOfStream() throws Exception {
        String name = "distrlog-mark-end-failure";
        DistributedLogManager dlm = createNewDLM(conf, name);

        long txid = 1;
        txid = writeAndMarkEndOfStream(dlm, txid);

        assertEquals(txid - 1, dlm.getLastTxId());
        LogRecord last = dlm.getLastLogRecord();
        assertEquals(txid - 1, last.getTransactionId());
        DLMTestUtil.verifyLogRecord(last);
        assert(dlm.isEndOfStreamMarked());

        LogWriter writer = null;
        boolean exceptionEncountered = false;
        try {
            writer = dlm.startLogSegmentNonPartitioned();
            for (long j = 1; j <= DEFAULT_SEGMENT_SIZE / 2; j++) {
                writer.write(DLMTestUtil.getLogRecordInstance(txid++));
            }
        } catch (EndOfStreamException exc) {
            exceptionEncountered = true;
        }
        writer.close();
        assert(exceptionEncountered);
    }

    @Test(timeout = 60000)
    public void testMarkEndOfStreamOnEmptyStream() throws Exception {
        markEndOfStreamOnEmptyLogSegment(0);
    }

    @Test(timeout = 60000)
    public void testMarkEndOfStreamOnClosedStream() throws Exception {
        markEndOfStreamOnEmptyLogSegment(3);
    }

    private void markEndOfStreamOnEmptyLogSegment(int numCompletedSegments) throws Exception {
        String name = "distrlog-mark-end-empty-" + numCompletedSegments;

        DistributedLogManager dlm = createNewDLM(conf, name);
        DLMTestUtil.generateCompletedLogSegments(dlm, conf, numCompletedSegments, DEFAULT_SEGMENT_SIZE);

        BKSyncLogWriter writer = (BKSyncLogWriter)dlm.startLogSegmentNonPartitioned();
        writer.markEndOfStream();

        LogReader reader = dlm.getInputStream(1);
        long numTrans = 0;
        boolean exceptionEncountered = false;
        try {
            LogRecord record = reader.readNext(false);
            long lastTxId = -1;
