/*
 *
 * 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 org.apache.bookkeeper.client.BookKeeperClientStats.CLIENT_SCOPE;
import static org.apache.bookkeeper.client.BookKeeperClientStats.SPECULATIVE_READ_COUNT;
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 java.util.BitSet;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import org.apache.bookkeeper.bookie.LocalBookieEnsemblePlacementPolicy;
import org.apache.bookkeeper.client.AsyncCallback.ReadCallback;
import org.apache.bookkeeper.client.BookKeeper.DigestType;
import org.apache.bookkeeper.conf.ClientConfiguration;
import org.apache.bookkeeper.net.BookieId;
import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
import org.apache.bookkeeper.test.TestStatsProvider;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This unit test tests ledger fencing.
 *
 */
public class TestSpeculativeBatchRead extends BookKeeperClusterTestCase {
    private static final Logger LOG = LoggerFactory.getLogger(TestSpeculativeBatchRead.class);

    private final DigestType digestType;
    byte[] passwd = "specPW".getBytes();

    public TestSpeculativeBatchRead() {
        super(10);
        this.digestType = DigestType.CRC32;
    }

    long getLedgerToRead(int ensemble, int quorum) throws Exception {
        byte[] data = "Data for test".getBytes();
        LedgerHandle l = bkc.createLedger(ensemble, quorum, digestType, passwd);
        for (int i = 0; i < 10; i++) {
            l.addEntry(data);
        }
        l.close();

        return l.getId();
    }

    @SuppressWarnings("deprecation")
    BookKeeperTestClient createClient(int specTimeout) throws Exception {
        ClientConfiguration conf = new ClientConfiguration()
            .setSpeculativeReadTimeout(specTimeout)
            .setReadTimeout(30000)
            .setUseV2WireProtocol(true)
            .setReorderReadSequenceEnabled(true)
            .setEnsemblePlacementPolicySlowBookies(true)
            .setMetadataServiceUri(zkUtil.getMetadataServiceUri());
        return new BookKeeperTestClient(conf, new TestStatsProvider());
    }

    class LatchCallback implements ReadCallback {
        CountDownLatch l = new CountDownLatch(1);
        boolean success = false;
        long startMillis = System.currentTimeMillis();
        long endMillis = Long.MAX_VALUE;

        public void readComplete(int rc,
                                 LedgerHandle lh,
                                 Enumeration<LedgerEntry> seq,
                                 Object ctx) {
            endMillis = System.currentTimeMillis();
            if (LOG.isDebugEnabled()) {
                LOG.debug("Got response {} {}", rc, getDuration());
            }
            success = rc == BKException.Code.OK;
            l.countDown();
        }

        long getDuration() {
            return endMillis - startMillis;
        }

        void expectSuccess(int milliseconds) throws Exception {
            boolean await = l.await(milliseconds, TimeUnit.MILLISECONDS);
            System.out.println(await);
        }

        void expectFail(int milliseconds) throws Exception {
            assertTrue(l.await(milliseconds, TimeUnit.MILLISECONDS));
            assertFalse(success);
        }

        void expectTimeout(int milliseconds) throws Exception {
            assertFalse(l.await(milliseconds, TimeUnit.MILLISECONDS));
        }
    }

    /**
     * Test basic speculative functionality.
     * - Create 2 clients with read timeout disabled, one with spec
     *   read enabled, the other not.
     * - create ledger
     * - sleep second bookie in ensemble
     * - read first entry, both should find on first bookie.
     * - read second bookie, spec client should find on bookie three,
     *   non spec client should hang.
     */
    @Test
    public void testSpeculativeRead() throws Exception {
        long id = getLedgerToRead(3, 2);
        BookKeeperTestClient bknospec = createClient(0); // disabled
        BookKeeperTestClient bkspec = createClient(2000);

        LedgerHandle lnospec = bknospec.openLedger(id, digestType, passwd);
        LedgerHandle lspec = bkspec.openLedger(id, digestType, passwd);

        // sleep second bookie
        CountDownLatch sleepLatch = new CountDownLatch(1);
        BookieId second = lnospec.getLedgerMetadata().getAllEnsembles().get(0L).get(1);
        sleepBookie(second, sleepLatch);

        try {
            // read first entry, both go to first bookie, should be fine
            LatchCallback nospeccb = new LatchCallback();
            LatchCallback speccb = new LatchCallback();
            lnospec.asyncBatchReadEntries(0, 1, 1024, nospeccb, null);
            lspec.asyncBatchReadEntries(0, 1, 1024, speccb, null);
            nospeccb.expectSuccess(2000);
            speccb.expectSuccess(2000);

            // read second entry, both look for second book, spec read client
            // tries third bookie, nonspec client hangs as read timeout is very long.
            nospeccb = new LatchCallback();
            speccb = new LatchCallback();
            lnospec.asyncReadEntries(1, 1, nospeccb, null);
            lspec.asyncReadEntries(1, 1, speccb, null);
            speccb.expectSuccess(4000);
            nospeccb.expectTimeout(4000);
            // Check that the second bookie is registered as slow at entryId 1
            RackawareEnsemblePlacementPolicy rep = (RackawareEnsemblePlacementPolicy) bkspec.getPlacementPolicy();
            assertTrue(rep.slowBookies.asMap().size() == 1);

            assertTrue(
                    "Stats should not reflect speculative reads if disabled",
                    bknospec.getTestStatsProvider()
                            .getCounter(CLIENT_SCOPE + "." + SPECULATIVE_READ_COUNT).get() == 0);
            assertTrue(
                    "Stats should reflect speculative reads",
                    bkspec.getTestStatsProvider()
                            .getCounter(CLIENT_SCOPE + "." + SPECULATIVE_READ_COUNT).get() > 0);
        } finally {
            sleepLatch.countDown();
            lspec.close();
            lnospec.close();
            bkspec.close();
            bknospec.close();
        }
    }

    /**
     * Test that if more than one replica is down, we can still read, as long as the quorum
     * size is larger than the number of down replicas.
     */
    @Test
    public void testSpeculativeReadMultipleReplicasDown() throws Exception {
        long id = getLedgerToRead(5, 5);
        int timeout = 5000;
        BookKeeper bkspec = createClient(timeout);

        LedgerHandle l = bkspec.openLedger(id, digestType, passwd);

        // sleep bookie 1, 2 & 4
        CountDownLatch sleepLatch = new CountDownLatch(1);
        sleepBookie(l.getLedgerMetadata().getAllEnsembles().get(0L).get(1), sleepLatch);
        sleepBookie(l.getLedgerMetadata().getAllEnsembles().get(0L).get(2), sleepLatch);
        sleepBookie(l.getLedgerMetadata().getAllEnsembles().get(0L).get(4), sleepLatch);

        try {
            // read first entry, should complete faster than timeout
            // as bookie 0 has the entry
            LatchCallback latch0 = new LatchCallback();
            l.asyncBatchReadEntries(0, 1, 1024, latch0, null);
            latch0.expectSuccess(timeout / 2);

            // second should have to hit two timeouts (bookie 1 & 2)
            // bookie 3 has the entry
            LatchCallback latch1 = new LatchCallback();
            l.asyncBatchReadEntries(1, 1, 1024, latch1, null);
            latch1.expectTimeout(timeout);
            latch1.expectSuccess(timeout * 2);
            LOG.info("Timeout {} latch1 duration {}", timeout, latch1.getDuration());
            assertTrue("should have taken longer than two timeouts, but less than 3",
                       latch1.getDuration() >= timeout * 2
                       && latch1.getDuration() < timeout * 3);

            // bookies 1 & 2 should be registered as slow bookies because of speculative reads
            Set<BookieId> expectedSlowBookies = new HashSet<>();
            expectedSlowBookies.add(l.getLedgerMetadata().getAllEnsembles().get(0L).get(1));
            expectedSlowBookies.add(l.getLedgerMetadata().getAllEnsembles().get(0L).get(2));
            assertEquals(((RackawareEnsemblePlacementPolicy) bkspec.getPlacementPolicy()).slowBookies.asMap().keySet(),
                expectedSlowBookies);

            // third should not hit timeouts since bookies 1 & 2 are registered as slow
            // bookie 3 has the entry
            LatchCallback latch2 = new LatchCallback();
            l.asyncBatchReadEntries(2, 1, 1024, latch2, null);
            latch2.expectSuccess(timeout);

            // fourth should have no timeout
            // bookie 3 has the entry
            LatchCallback latch3 = new LatchCallback();
            l.asyncBatchReadEntries(3, 1, 1024, latch3, null);
            latch3.expectSuccess(timeout / 2);

            // fifth should hit one timeout, (bookie 4)
            // bookie 0 has the entry
            LatchCallback latch4 = new LatchCallback();
            l.asyncBatchReadEntries(4, 1, 1024, latch4, null);
            latch4.expectTimeout(timeout / 2);
            latch4.expectSuccess(timeout);
            LOG.info("Timeout {} latch4 duration {}", timeout, latch4.getDuration());
            assertTrue("should have taken longer than one timeout, but less than 2",
                       latch4.getDuration() >= timeout
                       && latch4.getDuration() < timeout * 2);
        } finally {
            sleepLatch.countDown();
            l.close();
            bkspec.close();
        }
    }

    /**
     * Test that if after a speculative read is kicked off, the original read completes
     * nothing bad happens.
     */
    @Test
    public void testSpeculativeReadFirstReadCompleteIsOk() throws Exception {
        long id = getLedgerToRead(2, 2);
        int timeout = 1000;
        BookKeeper bkspec = createClient(timeout);

        LedgerHandle l = bkspec.openLedger(id, digestType, passwd);

        // sleep bookies
        CountDownLatch sleepLatch0 = new CountDownLatch(1);
        CountDownLatch sleepLatch1 = new CountDownLatch(1);
        sleepBookie(l.getLedgerMetadata().getAllEnsembles().get(0L).get(0), sleepLatch0);
        sleepBookie(l.getLedgerMetadata().getAllEnsembles().get(0L).get(1), sleepLatch1);

        try {
            // read goes to first bookie, spec read timeout occurs,
            // goes to second
            LatchCallback latch0 = new LatchCallback();
            l.asyncBatchReadEntries(0, 1, 1024, latch0, null);
            latch0.expectTimeout(timeout);

            // wake up first bookie
            sleepLatch0.countDown();
            latch0.expectSuccess(timeout / 2);

            sleepLatch1.countDown();

            // check we can read next entry without issue
            LatchCallback latch1 = new LatchCallback();
            l.asyncBatchReadEntries(1, 1, 1024, latch1, null);
            latch1.expectSuccess(timeout / 2);
        } finally {
            sleepLatch0.countDown();
            sleepLatch1.countDown();
            l.close();
            bkspec.close();
        }
    }

    /**
     * Unit test to check if the scheduled speculative task gets cancelled
     * on successful read.
     */
    @Test
    public void testSpeculativeReadScheduledTaskCancel() throws Exception {
        long id = getLedgerToRead(3, 2);
        int timeout = 1000;
        BookKeeper bkspec = createClient(timeout);
        LedgerHandle l = bkspec.openLedger(id, digestType, passwd);
        BatchedReadOp op = null;
        try {
            op = new BatchedReadOp(l, bkspec.getClientCtx(), 0, 5, 5120, false);
            op.initiate();
            op.future().get();
        } finally {
            assertNull("Speculative Read tasks must be null", op.getSpeculativeTask());
        }
    }

    /**
     * Unit test for the speculative read scheduling method.
     */
    @Test
    public void testSpeculativeReadScheduling() throws Exception {
        long id = getLedgerToRead(3, 2);
        int timeout = 1000;
        BookKeeper bkspec = createClient(timeout);

        LedgerHandle l = bkspec.openLedger(id, digestType, passwd);

        List<BookieId> ensemble = l.getLedgerMetadata().getAllEnsembles().get(0L);
        BitSet allHosts = new BitSet(ensemble.size());
        for (int i = 0; i < ensemble.size(); i++) {
            allHosts.set(i, true);
        }
        BitSet noHost = new BitSet(ensemble.size());
        BitSet secondHostOnly = new BitSet(ensemble.size());
        secondHostOnly.set(1, true);
        BatchedReadOp.LedgerEntryRequest req0 = null, req2 = null, req4 = null;
        try {
            BatchedReadOp op = new BatchedReadOp(l, bkspec.getClientCtx(), 0, 5, 5120, false);
            // if we've already heard from all hosts,
            // we only send the initial read
            req0 = op.new SequenceReadRequest(ensemble, l.getId(), 0, 1, 1024);
            assertTrue("Should have sent to first",
                       req0.maybeSendSpeculativeRead(allHosts).equals(ensemble.get(0)));
            assertNull("Should not have sent another",
                       req0.maybeSendSpeculativeRead(allHosts));

            // if we have heard from some hosts, but not one we have sent to
            // send again
            req2 = op.new SequenceReadRequest(ensemble, l.getId(), 2, 1, 1024);
            assertTrue("Should have sent to third",
                       req2.maybeSendSpeculativeRead(noHost).equals(ensemble.get(2)));
            assertTrue("Should have sent to first",
                       req2.maybeSendSpeculativeRead(secondHostOnly).equals(ensemble.get(0)));

            // if we have heard from some hosts, which includes one we sent to
            // do not read again
            req4 = op.new SequenceReadRequest(ensemble, l.getId(), 4, 1, 1024);
            assertTrue("Should have sent to second",
                       req4.maybeSendSpeculativeRead(noHost).equals(ensemble.get(1)));
            assertNull("Should not have sent another",
                       req4.maybeSendSpeculativeRead(secondHostOnly));
        } finally {
            for (BatchedReadOp.LedgerEntryRequest req
                     : new BatchedReadOp.LedgerEntryRequest[] { req0, req2, req4 }) {
                if (req != null) {
                    int i = 0;
                    while (!req.isComplete()) {
                        if (i++ > 10) {
                            break; // wait for up to 10 seconds
                        }
                        Thread.sleep(1000);
                    }
                    assertTrue("Request should be done", req.isComplete());
                }
            }

            l.close();
            bkspec.close();
        }
    }

    @Test
    public void testSequenceReadLocalEnsemble() throws Exception {
        ClientConfiguration conf = new ClientConfiguration()
                .setSpeculativeReadTimeout(1000)
                .setEnsemblePlacementPolicy(LocalBookieEnsemblePlacementPolicy.class)
                .setReorderReadSequenceEnabled(true)
                .setEnsemblePlacementPolicySlowBookies(true)
                .setMetadataServiceUri(zkUtil.getMetadataServiceUri());
        BookKeeper bkspec = new BookKeeperTestClient(conf, new TestStatsProvider());
        LedgerHandle l = bkspec.createLedger(1, 1, digestType, passwd);
        List<BookieId> ensemble = l.getLedgerMetadata().getAllEnsembles().get(0L);
        BatchedReadOp op = new BatchedReadOp(l, bkspec.getClientCtx(), 0, 5, 5120, false);
        BatchedReadOp.LedgerEntryRequest req0 = op.new SequenceReadRequest(ensemble, l.getId(), 0, 1, 1024);
        assertNotNull(req0.writeSet);
    }
}
