content: /*
   *
   * 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.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertTrue;
  
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ThreadLocalRandom;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicInteger;
  import org.apache.bookkeeper.conf.ClientConfiguration;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.proto.BookieClientImpl;
  import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.GenericCallback;
  import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
  import org.awaitility.Awaitility;
  import org.junit.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Test a slow bookie.
   */
  @SuppressWarnings("deprecation")
  public class SlowBookieTest extends BookKeeperClusterTestCase {
      private static final Logger LOG = LoggerFactory.getLogger(SlowBookieTest.class);
  
      final byte[] entry = "Test Entry".getBytes();
  
      public SlowBookieTest() {
          super(4);
          baseConf.setNumAddWorkerThreads(0);
          baseConf.setNumReadWorkerThreads(0);
      }
  
      @Test
      public void testSlowBookie() throws Exception {
          ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(360)
              .setMetadataServiceUri(zkUtil.getMetadataServiceUri());
  
          BookKeeper bkc = new BookKeeper(conf);
  
          LedgerHandle lh = bkc.createLedger(4, 3, 2, BookKeeper.DigestType.CRC32, new byte[] {});
  
          byte[] entry = "Test Entry".getBytes();
          for (int i = 0; i < 10; i++) {
              lh.addEntry(entry);
          }
          final CountDownLatch b0latch = new CountDownLatch(1);
          final CountDownLatch b1latch = new CountDownLatch(1);
          final CountDownLatch addEntrylatch = new CountDownLatch(1);
          List<BookieId> curEns = lh.getCurrentEnsemble();
          try {
              sleepBookie(curEns.get(0), b0latch);
              for (int i = 0; i < 10; i++) {
                  lh.addEntry(entry);
              }
              sleepBookie(curEns.get(2), b1latch); // should cover all quorums
  
              final AtomicInteger i = new AtomicInteger(0xdeadbeef);
              AsyncCallback.AddCallback cb = new AsyncCallback.AddCallback() {
                      public void addComplete(int rc, LedgerHandle lh, long entryId, Object ctx) {
                          i.set(rc);
                          addEntrylatch.countDown();
                      }
                  };
              lh.asyncAddEntry(entry, cb, null);
  
              Awaitility.await().untilAsserted(() ->
                  assertEquals("Successfully added entry!", 0xdeadbeef, i.get()));
              b0latch.countDown();
              b1latch.countDown();
              addEntrylatch.await(4000, TimeUnit.MILLISECONDS);
              assertEquals("Failed to add entry!", BKException.Code.OK, i.get());
          } finally {
              b0latch.countDown();
              b1latch.countDown();
              addEntrylatch.countDown();
          }
      }
  
      @Test
      public void testBookieFailureWithSlowBookie() throws Exception {
          ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(5)
              .setMetadataServiceUri(zkUtil.getMetadataServiceUri());
  
          BookKeeper bkc = new BookKeeper(conf);
  
          byte[] pwd = new byte[] {};
          final LedgerHandle lh = bkc.createLedger(4, 3, 2, BookKeeper.DigestType.CRC32, pwd);
          final AtomicBoolean finished = new AtomicBoolean(false);
          final AtomicBoolean failTest = new AtomicBoolean(false);
          Thread t = new Thread() {
                  public void run() {
                      try {
                          while (!finished.get()) {
                              lh.addEntry(entry);
                          }
                      } catch (Exception e) {
                          LOG.error("Exception in add entry thread", e);
                          failTest.set(true);
                      }
                  }
              };
          t.start();
          final CountDownLatch b0latch = new CountDownLatch(1);
  
          startNewBookie();
          sleepBookie(getBookie(0), b0latch);
  
          Thread.sleep(10000);
          b0latch.countDown();
  
          finished.set(true);
          t.join();
  
          assertFalse(failTest.get());
  
          lh.close();
  
          LedgerHandle lh2 = bkc.openLedger(lh.getId(), BookKeeper.DigestType.CRC32, pwd);
          LedgerChecker lc = new LedgerChecker(bkc);
          final CountDownLatch checklatch = new CountDownLatch(1);
          final AtomicInteger numFragments = new AtomicInteger(-1);
          lc.checkLedger(lh2, new GenericCallback<Set<LedgerFragment>>() {
                  public void operationComplete(int rc, Set<LedgerFragment> badFragments) {
                      if (LOG.isDebugEnabled()) {
                          LOG.debug("Checked ledgers returned {} {}", rc, badFragments);
                      }
                      if (rc == BKException.Code.OK) {
                          numFragments.set(badFragments.size());
                      }
                      checklatch.countDown();
                  }
              });
          checklatch.await();
          assertEquals("There should be no missing fragments", 0, numFragments.get());
      }
  
      @Test
      public void testSlowBookieAndBackpressureOn() throws Exception {
          final ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(5)
                  .setAddEntryTimeout(1)
                  .setAddEntryQuorumTimeout(1)
                  .setNumChannelsPerBookie(1)
                  .setZkServers(zkUtil.getZooKeeperConnectString())
                  .setClientWriteBufferLowWaterMark(1)
                  .setClientWriteBufferHighWaterMark(entry.length - 1)
                  .setWaitTimeoutOnBackpressureMillis(5000);
  
          final boolean expectWriteError = false;
          final boolean expectFailedTest = false;
  
          try (LedgerHandle lh = doBackPressureTest(entry, conf, expectWriteError, expectFailedTest, 2000)) {
              assertTrue(lh.readLastConfirmed() < 5);
          }
      }
  
      @Test
      public void testSlowBookieAndFastFailOn() throws Exception {
          final ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(5)
                  .setAddEntryTimeout(1)
                  .setAddEntryQuorumTimeout(1)
                  .setNumChannelsPerBookie(1)
                  .setZkServers(zkUtil.getZooKeeperConnectString())
                  .setClientWriteBufferLowWaterMark(1)
                  .setClientWriteBufferHighWaterMark(2)
                  .setWaitTimeoutOnBackpressureMillis(0);
  
          final boolean expectWriteError = true;
          final boolean expectFailedTest = false;
  
          try (LedgerHandle lh = doBackPressureTest(entry, conf, expectWriteError, expectFailedTest, 1000)) {
              assertTrue(lh.readLastConfirmed() < 5);
          }
      }
  
      @Test
      public void testSlowBookieAndNoBackpressure() throws Exception {
          final ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(5)
                  .setAddEntryTimeout(1)
                  .setAddEntryQuorumTimeout(1)
                  .setNumChannelsPerBookie(1)
                  .setZkServers(zkUtil.getZooKeeperConnectString())
                  .setClientWriteBufferLowWaterMark(1)
                  .setClientWriteBufferHighWaterMark(entry.length - 1)
                  .setWaitTimeoutOnBackpressureMillis(-1);
  
          final boolean expectWriteError = false;
          final boolean expectFailedTest = false;
  
          try (LedgerHandle lh = doBackPressureTest(entry, conf, expectWriteError, expectFailedTest, 4000)) {
              assertTrue(lh.readLastConfirmed() > 90);
          }
      }
  
      private LedgerHandle doBackPressureTest(byte[] entry, ClientConfiguration conf,
                                              boolean expectWriteError, boolean expectFailedTest,
                                              long sleepInMillis) throws Exception {
          BookKeeper bkc = new BookKeeper(conf);
  
          byte[] pwd = new byte[] {};
          final LedgerHandle lh = bkc.createLedger(4, 3, 1, BookKeeper.DigestType.CRC32, pwd);
          lh.addEntry(entry);
  
          final AtomicBoolean finished = new AtomicBoolean(false);
          final AtomicBoolean failTest = new AtomicBoolean(false);
          final AtomicBoolean writeError = new AtomicBoolean(false);
          Thread t = new Thread(() -> {
              try {
                  int count = 0;
                  while (!finished.get()) {
                      lh.asyncAddEntry(entry, (rc, lh1, entryId, ctx) -> {
                          if (rc != BKException.Code.OK) {
                              finished.set(true);
                              writeError.set(true);
                          }
                      }, null);
                      if (++count > 100) {
                          finished.set(true);
                      }
                  }
              } catch (Exception e) {
                  LOG.error("Exception in add entry thread", e);
                  failTest.set(true);
              }
          });
          final CountDownLatch b0latch = new CountDownLatch(1);
          final CountDownLatch b0latch2 = new CountDownLatch(1);
  
  
          sleepBookie(getBookie(0), b0latch);
          sleepBookie(getBookie(1), b0latch2);
  
          setTargetChannelState(bkc, getBookie(0), 0, false);
          setTargetChannelState(bkc, getBookie(1), 0, false);
  
          t.start();
  
          Thread.sleep(sleepInMillis);
  
          finished.set(true);
  
          b0latch.countDown();
          b0latch2.countDown();
          setTargetChannelState(bkc, getBookie(0), 0, true);
          setTargetChannelState(bkc, getBookie(1), 0, true);
  
          t.join();
  
          assertEquals("write error", expectWriteError, writeError.get());
          assertEquals("test failure", expectFailedTest, failTest.get());
  
          lh.close();
  
          LedgerHandle lh2 = bkc.openLedger(lh.getId(), BookKeeper.DigestType.CRC32, pwd);
          LedgerChecker lc = new LedgerChecker(bkc);
          final CountDownLatch checkLatch = new CountDownLatch(1);
          final AtomicInteger numFragments = new AtomicInteger(-1);
          lc.checkLedger(lh2, (rc, fragments) -> {
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Checked ledgers returned {} {}", rc, fragments);
              }
              if (rc == BKException.Code.OK) {
                  numFragments.set(fragments.size());
                  LOG.error("Checked ledgers returned {} {}", rc, fragments);
              }
              checkLatch.countDown();
          });
          checkLatch.await();
          assertEquals("There should be no missing fragments", 0, numFragments.get());
  
          return lh2;
      }
  
      private void setTargetChannelState(BookKeeper bkc, BookieId address,
                                         long key, boolean writable) throws Exception {
          ((BookieClientImpl) bkc.getBookieClient()).lookupClient(address).obtain((rc, pcbc) -> {
              pcbc.setWritable(writable);
          }, key);
      }
  
      @Test
      public void testWriteSetWriteableCheck() throws Exception {
          final ClientConfiguration conf = new ClientConfiguration();
          conf.setMetadataServiceUri(zkUtil.getMetadataServiceUri());
          BookKeeper bkc = new BookKeeper(conf);
  
          byte[] pwd = new byte[]{};
          try (LedgerHandle lh = bkc.createLedger(4, 2, 2, BookKeeper.DigestType.CRC32, pwd)) {
              lh.addEntry(entry); // [b0, b1]
              long entryId = lh.addEntry(entry); // [b1, b2]
  
              long nextEntryId = entryId + 1;
              RoundRobinDistributionSchedule schedule = new RoundRobinDistributionSchedule(2, 2, 4);
              DistributionSchedule.WriteSet writeSet = schedule.getWriteSet(nextEntryId);
  
              // b2 or b3 is no more writeable
              int slowBookieIndex = writeSet.get(ThreadLocalRandom.current().nextInt(writeSet.size()));
              List<BookieId> curEns = lh.getCurrentEnsemble();
  
              // Trigger connection to the bookie service first
              bkc.getBookieInfo().get(curEns.get(slowBookieIndex));
              // then mock channel is not writable
              setTargetChannelState(bkc, curEns.get(slowBookieIndex), 0, false);
  
              boolean isWriteable = lh.waitForWritable(writeSet, 0, 1000);
              assertFalse("We should check b2,b3 both are not writeable", isWriteable);
          }
      }
  
      @Test
      public void testManyBookieFailureWithSlowBookies() throws Exception {
          ClientConfiguration conf = new ClientConfiguration();
          conf.setReadTimeout(5)
              .setMetadataServiceUri(zkUtil.getMetadataServiceUri());
  
          BookKeeper bkc = new BookKeeper(conf);
  
          byte[] pwd = new byte[] {};
          final LedgerHandle lh = bkc.createLedger(4, 3, 2, BookKeeper.DigestType.CRC32, pwd);
          final AtomicBoolean finished = new AtomicBoolean(false);
          final AtomicBoolean failTest = new AtomicBoolean(false);
          Thread t = new Thread() {
                  public void run() {
                      try {
                          while (!finished.get()) {
                              lh.addEntry(entry);
                              Thread.sleep(1);
                          }
                      } catch (Exception e) {
                          LOG.error("Exception in add entry thread", e);
                          failTest.set(true);
                      }
                  }
              };
          t.start();
          final CountDownLatch b0latch = new CountDownLatch(1);
          final CountDownLatch b1latch = new CountDownLatch(1);
  
          startNewBookie();
          startNewBookie();
  
          sleepBookie(getBookie(0), b0latch);
          sleepBookie(getBookie(1), b1latch);
  
          Thread.sleep(10000);
          b0latch.countDown();
          b1latch.countDown();
          finished.set(true);
          t.join();
  
          assertFalse(failTest.get());
  
          lh.close();
  
          LedgerHandle lh2 = bkc.openLedger(lh.getId(), BookKeeper.DigestType.CRC32, pwd);
          LedgerChecker lc = new LedgerChecker(bkc);
          final CountDownLatch checklatch = new CountDownLatch(1);
          final AtomicInteger numFragments = new AtomicInteger(-1);
          lc.checkLedger(lh2, new GenericCallback<Set<LedgerFragment>>() {
                  public void operationComplete(int rc, Set<LedgerFragment> fragments) {
                      if (LOG.isDebugEnabled()) {
                          LOG.debug("Checked ledgers returned {} {}", rc, fragments);
                      }
                      if (rc == BKException.Code.OK) {
                          numFragments.set(fragments.size());
                      }
                      checklatch.countDown();
                  }
              });
          checklatch.await();
          assertEquals("There should be no missing fragments", 0, numFragments.get());
      }
  
      @Test
      public void testWaitForWritable() throws Exception {
          final ClientConfiguration conf = new ClientConfiguration();
          conf.setMetadataServiceUri(zkUtil.getMetadataServiceUri());
          BookKeeper bkc = new BookKeeper(conf);
  
          byte[] pwd = new byte[]{};
          try (LedgerHandle lh = bkc.createLedger(1, 1, 1, BookKeeper.DigestType.CRC32, pwd)) {
              long entryId = lh.addEntry(this.entry);
  
              RoundRobinDistributionSchedule schedule = new RoundRobinDistributionSchedule(1, 1, 1);
              DistributionSchedule.WriteSet writeSet = schedule.getWriteSet(entryId);
  
              int slowBookieIndex = writeSet.get(ThreadLocalRandom.current().nextInt(writeSet.size()));
              List<BookieId> curEns = lh.getCurrentEnsemble();
  
              // disable channel writable
              setTargetChannelState(bkc, curEns.get(slowBookieIndex), 0, false);
  
              AtomicBoolean isWriteable = new AtomicBoolean(false);
              final long timeout = 10000;
  
              // waitForWritable async
              new Thread(() -> isWriteable.set(lh.waitForWritable(writeSet, 0, timeout))).start();
  
              Awaitility.await().pollDelay(5, TimeUnit.SECONDS).untilAsserted(() -> assertFalse(isWriteable.get()));
  
              // enable channel writable
              setTargetChannelState(bkc, curEns.get(slowBookieIndex), 0, true);
              Awaitility.await().untilAsserted(() -> assertTrue(isWriteable.get()));
          }
      }
  
  }
  