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.metadata.etcd;
  
  import static org.apache.bookkeeper.common.concurrent.FutureUtils.result;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotEquals;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import com.google.common.collect.Sets;
  import io.etcd.jetcd.Client;
  import java.util.ArrayList;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.BrokenBarrierException;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CyclicBarrier;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.atomic.AtomicInteger;
  import lombok.Cleanup;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.bookie.BookieException;
  import org.apache.bookkeeper.bookie.BookieException.MetadataStoreException;
  import org.apache.bookkeeper.common.concurrent.FutureUtils;
  import org.apache.bookkeeper.discover.BookieServiceInfo;
  import org.apache.bookkeeper.discover.RegistrationClient;
  import org.apache.bookkeeper.discover.RegistrationClient.RegistrationListener;
  import org.apache.bookkeeper.metadata.etcd.testing.EtcdTestBase;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.versioning.LongVersion;
  import org.apache.bookkeeper.versioning.Version;
  import org.apache.bookkeeper.versioning.Version.Occurred;
  import org.apache.bookkeeper.versioning.Versioned;
  import org.apache.commons.lang3.RandomStringUtils;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.TestName;
  
  /**
   * Test etcd based bookie registration.
   */
  @Slf4j
  public class EtcdRegistrationTest extends EtcdTestBase {
  
      static BookieId newBookie(int i) {
          return BookieId.parse("127.0.0.1:" + (3181 + i));
      }
  
      @Rule
      public final TestName runtime = new TestName();
  
      private String scope;
      private RegistrationClient regClient;
  
      protected static RegistrationListener newRegistrationListener(
          LinkedBlockingQueue<Versioned<Set<BookieId>>> notifications) {
          return bookies -> {
              log.info("Received new bookies: {}", bookies);
              try {
                  notifications.put(bookies);
              } catch (InterruptedException e) {
                  log.error("Interrupted at enqueuing updated key set", e);
              }
          };
      }
  
      @Before
      @Override
      public void setUp() throws Exception {
          super.setUp();
          this.scope = RandomStringUtils.randomAlphabetic(16);
          this.regClient = new EtcdRegistrationClient(scope, etcdClient);
      }
  
      @After
      @Override
      public void tearDown() throws Exception {
          this.regClient.close();
          super.tearDown();
      }
  
      interface MultiBookiesTester {
  
          void test(String scope, int numBookies, boolean readonly) throws Exception;
  
      }
  
      private static void runNumBookiesTest(final String scope,
                                            final int numBookies,
                                            final boolean readonly,
                                            MultiBookiesTester tester) throws Exception {
  
          final List<EtcdRegistrationManager> bookies = createNumBookies(readonly, numBookies, scope);
          try {
              tester.test(scope, numBookies, readonly);
          } finally {
              bookies.forEach(EtcdRegistrationManager::close);
          }
  
      }
  
      @Test
      public void testRegisterWritableBookies() throws Exception {
          testRegisterBookie(false);
      }
  
      @Test
      public void testRegisterReadonlyBookies() throws Exception {
          testRegisterBookie(true);
      }
  
      private void testRegisterBookie(boolean readonly) throws Exception {
          runNumBookiesTest(scope, 3, readonly, (scope, numBookies, ro) -> {
              Set<BookieId> expectedBookies = Sets.newHashSet();
              for (int i = 0; i < numBookies; i++) {
                  expectedBookies.add(newBookie(i));
              }
              Set<BookieId> writableBookies = result(regClient.getWritableBookies()).getValue();
              Set<BookieId> readonlyBookies = result(regClient.getReadOnlyBookies()).getValue();
              if (ro) {
                  assertEquals(0, writableBookies.size());
                  assertEquals(numBookies, readonlyBookies.size());
                  assertEquals(expectedBookies, readonlyBookies);
              } else {
                  assertEquals(0, readonlyBookies.size());
                  assertEquals(numBookies, writableBookies.size());
                  assertEquals(expectedBookies, writableBookies);
              }
  
          });
      }
  
      @Test
      public void testWatchWritableBookies() throws Exception {
          testWatchBookies(false);
      }
  
      @Test
      public void testWatchReadonlyBookies() throws Exception {
          testWatchBookies(true);
      }
  
      private void testWatchBookies(boolean readonly) throws Exception {
          LinkedBlockingQueue<Versioned<Set<BookieId>>> writableChanges = new LinkedBlockingQueue<>();
          LinkedBlockingQueue<Versioned<Set<BookieId>>> readonlyChanges = new LinkedBlockingQueue<>();
          result(regClient.watchReadOnlyBookies(newRegistrationListener(readonlyChanges)));
          result(regClient.watchWritableBookies(newRegistrationListener(writableChanges)));
          Versioned<Set<BookieId>> versionedBookies = writableChanges.take();
          assertTrue(versionedBookies.getValue().isEmpty());
          versionedBookies = readonlyChanges.take();
          assertTrue(versionedBookies.getValue().isEmpty());
  
          final int numBookies = 3;
          final List<EtcdRegistrationManager> bookies = createNumBookies(readonly, numBookies, scope, 1);
  
          LinkedBlockingQueue<Versioned<Set<BookieId>>> changes;
          if (readonly) {
              changes = readonlyChanges;
          } else {
              changes = writableChanges;
          }
  
          Version preVersion = new LongVersion(-1);
          Set<BookieId> expectedBookies = new HashSet<>();
          for (int i = 0; i < numBookies; i++) {
              BookieId address = newBookie(i);
              expectedBookies.add(address);
  
              versionedBookies = changes.take();
              Version curVersion = versionedBookies.getVersion();
              assertEquals(Occurred.AFTER, curVersion.compare(preVersion));
              assertEquals(expectedBookies, versionedBookies.getValue());
              preVersion = curVersion;
          }
  
          bookies.forEach(EtcdRegistrationManager::close);
          for (int i = 0; i < numBookies; i++) {
              versionedBookies = changes.take();
              Version curVersion = versionedBookies.getVersion();
              assertEquals(Occurred.AFTER, curVersion.compare(preVersion));
              assertEquals(numBookies - i - 1, versionedBookies.getValue().size());
              preVersion = curVersion;
          }
          if (readonly) {
              assertEquals(0, writableChanges.size());
          } else {
              assertEquals(0, readonlyChanges.size());
          }
      }
  
      private static List<EtcdRegistrationManager> createNumBookies(boolean readonly,
                                                                    int numBookies,
                                                                    String scope,
                                                                    long ttlSeconds) throws BookieException {
          List<EtcdRegistrationManager> bookies = new ArrayList<>(numBookies);
          for (int i = 0; i < numBookies; i++) {
              Client client = newEtcdClient();
              EtcdRegistrationManager regMgr = new EtcdRegistrationManager(client, scope, ttlSeconds);
              bookies.add(regMgr);
              regMgr.registerBookie(newBookie(i), readonly, BookieServiceInfo.EMPTY);
          }
          return bookies;
      }
  
      private static List<EtcdRegistrationManager> createNumBookies(boolean readonly,
                                                                    int numBookies,
                                                                    String scope) throws BookieException {
          return createNumBookies(readonly, numBookies, scope, 60);
      }
  
      @Test
      public void testRegisterBookieWaitUntilPreviousExpiredSuccess() throws Exception {
          long ttlSeconds = 1;
          long leaseId = -0xabcd;
          BookieId bookieId = BookieId.parse(runtime.getMethodName() + ":3181");
          try (EtcdRegistrationManager regManager = new EtcdRegistrationManager(
              newEtcdClient(), scope, ttlSeconds)
          ) {
              regManager.registerBookie(bookieId, false, BookieServiceInfo.EMPTY);
              leaseId = regManager.getBkRegister().getLeaseId();
              log.info("Registered bookie under scope '{}' with lease = {}", scope, leaseId);
          }
          assertNotEquals(-0xabcd, leaseId);
          final long prevLeaseId = leaseId;
          try (EtcdRegistrationManager regManager = new EtcdRegistrationManager(
              newEtcdClient(), scope, 100000 * ttlSeconds)
          ) {
              regManager.registerBookie(bookieId, false, BookieServiceInfo.EMPTY);
              leaseId = regManager.getBkRegister().getLeaseId();
              log.info("Registered bookie under scope '{}' with new lease = {}", scope, leaseId);
          }
          assertNotEquals(prevLeaseId, leaseId);
      }
  
      @Test
      public void testRegisterBookieWaitUntilPreviousExpiredFailure() throws Exception {
          long ttlSeconds = 1;
          long leaseId = -0xabcd;
          BookieId bookieId = BookieId.parse(runtime.getMethodName() + ":3181");
          try (EtcdRegistrationManager regManager = new EtcdRegistrationManager(
              newEtcdClient(), scope, 10000000 * ttlSeconds)
          ) {
              regManager.registerBookie(bookieId, false, BookieServiceInfo.EMPTY);
              leaseId = regManager.getBkRegister().getLeaseId();
              log.info("Registered bookie under scope '{}' with lease = {}", scope, leaseId);
          }
          assertNotEquals(-0xabcd, leaseId);
          try (EtcdRegistrationManager regManager = new EtcdRegistrationManager(
              newEtcdClient(), scope,  ttlSeconds)
          ) {
              regManager.registerBookie(bookieId, false, BookieServiceInfo.EMPTY);
              fail("Should fail to register bookie under scope '{}'"
                  + " since previous registration has not been expired yet");
          } catch (MetadataStoreException mse) {
              log.info("Encountered exception on registering bookie under scope '{}'", scope, mse);
              // expected
          }
      }
  
      @Test
      public void testRegisterWritableBookieWithSameLeaseId() throws Exception {
          testRegisterBookieWithSameLeaseId(false);
      }
  
      @Test
      public void testRegisterReadonlyBookieWithSameLeaseId() throws Exception {
          testRegisterBookieWithSameLeaseId(true);
      }
  
      private void testRegisterBookieWithSameLeaseId(boolean readonly) throws Exception {
          long ttlSeconds = 1;
          long leaseId = -0xabcd;
          BookieId bookieId = BookieId.parse(runtime.getMethodName() + ":3181");
          try (EtcdRegistrationManager regManager = new EtcdRegistrationManager(
              newEtcdClient(), scope, 10000000 * ttlSeconds)
          ) {
              regManager.registerBookie(bookieId, readonly, BookieServiceInfo.EMPTY);
              leaseId = regManager.getBkRegister().getLeaseId();
              log.info("Registered bookie under scope '{}' with lease = {}", scope, leaseId);
              log.info("Trying to register using same lease '{}'", leaseId);
              try (EtcdRegistrationManager regManager2 = new EtcdRegistrationManager(
                  regManager.getClient(), scope, regManager.getBkRegister()
              )) {
                  regManager.registerBookie(bookieId, readonly, BookieServiceInfo.EMPTY);
              }
          }
      }
  
      private Set<BookieId> getBookies(boolean readonly) throws Exception {
          Set<BookieId> bookies;
          if (readonly) {
              bookies = result(regClient.getReadOnlyBookies()).getValue();
          } else {
              bookies = result(regClient.getWritableBookies()).getValue();
          }
          return bookies;
      }
  
      @Test
      public void testRegisterUnregisterWritableBookie() throws Exception {
          testRegisterUnregister(false);
      }
  
      @Test
      public void testRegisterUnregisterReadonlyBookie() throws Exception {
          testRegisterUnregister(true);
      }
  
      private void testRegisterUnregister(boolean readonly) throws Exception {
          String bookieIdStr = runtime.getMethodName();
          if (readonly) {
              bookieIdStr += "-readonly";
          }
          bookieIdStr += ":3181";
          BookieId bookieId = BookieId.parse(bookieIdStr);
          try (EtcdRegistrationManager regMgr = new EtcdRegistrationManager(
              newEtcdClient(), scope, 1000000000
          )) {
              // before registration
              Set<BookieId> bookies = getBookies(readonly);
              log.info("before registration : bookies = {}", bookies);
              assertEquals(0, bookies.size());
              // registered
              regMgr.registerBookie(bookieId, readonly, BookieServiceInfo.EMPTY);
              bookies = getBookies(readonly);
              log.info("after registered: bookies = {}", bookies);
              assertEquals(1, bookies.size());
              assertEquals(
                  Sets.newHashSet(bookieId),
                  bookies);
              // unregistered
              regMgr.unregisterBookie(bookieId, readonly);
              bookies = getBookies(readonly);
              log.info("after unregistered: bookies = {}", bookies);
              assertEquals(0, bookies.size());
          }
      }
  
      @Test
      public void testConcurrentWritableRegistration() throws Exception {
          testConcurrentRegistration(false);
      }
  
      @Test
      public void testConcurrentReadonlyRegistration() throws Exception {
          testConcurrentRegistration(true);
      }
  
      private void testConcurrentRegistration(boolean readonly) throws Exception {
          final BookieId bookieId;
          if (readonly) {
              bookieId = BookieId.parse(runtime.getMethodName() + "-readonly:3181");
          } else {
              bookieId = BookieId.parse(runtime.getMethodName() + ":3181");
          }
          final int numBookies = 10;
          @Cleanup("shutdown")
          ExecutorService executor = Executors.newFixedThreadPool(numBookies);
          final CyclicBarrier startBarrier = new CyclicBarrier(numBookies);
          final CyclicBarrier completeBarrier = new CyclicBarrier(numBookies);
          final CompletableFuture<Void> doneFuture = new CompletableFuture<>();
          final AtomicInteger numSuccesses = new AtomicInteger(0);
          final AtomicInteger numFailures = new AtomicInteger(0);
          for (int i = 0; i < numBookies; i++) {
              executor.submit(() -> {
                  try (EtcdRegistrationManager regMgr = new EtcdRegistrationManager(
                      newEtcdClient(), scope, 1
                  )) {
                      try {
                          startBarrier.await();
                          regMgr.registerBookie(bookieId, readonly, BookieServiceInfo.EMPTY);
                          numSuccesses.incrementAndGet();
                      } catch (InterruptedException e) {
                          log.warn("Interrupted at waiting for the other threads to start", e);
                      } catch (BrokenBarrierException e) {
                          log.warn("Start barrier is broken", e);
                      } catch (BookieException e) {
                          numFailures.incrementAndGet();
                      }
                      try {
                          completeBarrier.await();
                      } catch (InterruptedException e) {
                          log.warn("Interrupted at waiting for the other threads to complete", e);
                      } catch (BrokenBarrierException e) {
                          log.warn("Complete barrier is broken", e);
                      }
                      FutureUtils.complete(doneFuture, null);
                  }
              });
          }
          doneFuture.join();
          assertEquals(1, numSuccesses.get());
          assertEquals(numBookies - 1, numFailures.get());
      }
  
  }
  