io.atomix.primitive.session.SessionClient Java Examples

private CompletableFuture<Void> close(Function<SessionClient, CompletableFuture<Void>> closeFunction) {
  if (closeFuture == null) {
    synchronized (this) {
      if (closeFuture == null) {
        connected = false;
        SessionClient session = this.session;
        if (session != null) {
          closeFuture = closeFunction.apply(session);
        } else if (closeFuture != null) {
          closeFuture = connectFuture.thenCompose(closeFunction);
        } else {
          closeFuture = CompletableFuture.completedFuture(null);
        }
      }
    }
  }
  return closeFuture;
}
 

@Override
public <S> ProxyClient<S> newProxy(String primitiveName, PrimitiveType primitiveType, Class<S> serviceType, ServiceConfig serviceConfig, PartitionService partitionService) {
  PartitionGroup partitionGroup = partitionService.getPartitionGroup(this);
  if (partitionGroup == null) {
    throw new ConfigurationException("No Raft partition group matching the configured protocol exists");
  }

  Collection<SessionClient> partitions = partitionGroup.getPartitions().stream()
      .map(partition -> ((PrimaryBackupPartition) partition).getClient()
          .sessionBuilder(primitiveName, primitiveType, serviceConfig)
          .withConsistency(config.getConsistency())
          .withReplication(config.getReplication())
          .withRecovery(config.getRecovery())
          .withNumBackups(config.getBackups())
          .withMaxRetries(config.getMaxRetries())
          .withRetryDelay(config.getRetryDelay())
          .build())
      .collect(Collectors.toList());
  return new DefaultProxyClient<>(primitiveName, primitiveType, this, serviceType, partitions, config.getPartitioner());
}
 

/**
 * Recursively connects to the partition.
 */
private void connect(int attempt, CompletableFuture<SessionClient> future) {
  if (attempt > MAX_ATTEMPTS) {
    future.completeExceptionally(new PrimitiveException.Unavailable());
    return;
  }

  primaryElection.getTerm().whenCompleteAsync((term, error) -> {
    if (error == null) {
      if (term.primary() == null) {
        future.completeExceptionally(new PrimitiveException.Unavailable());
      } else {
        this.term = term;
        protocol.registerEventListener(sessionId, this::handleEvent, threadContext);
        future.complete(this);
      }
    } else {
      Throwable cause = Throwables.getRootCause(error);
      if (cause instanceof PrimitiveException.Unavailable || cause instanceof TimeoutException) {
        threadContext.schedule(Duration.ofMillis(RETRY_DELAY), () -> connect(attempt + 1, future));
      } else {
        future.completeExceptionally(new PrimitiveException.Unavailable());
      }
    }
  }, threadContext);
}
 

public PrimaryBackupSessionClient(
    String clientName,
    PartitionId partitionId,
    SessionId sessionId,
    PrimitiveType primitiveType,
    PrimitiveDescriptor descriptor,
    ClusterMembershipService clusterMembershipService,
    PrimaryBackupClientProtocol protocol,
    PrimaryElection primaryElection,
    ThreadContext threadContext) {
  this.partitionId = checkNotNull(partitionId);
  this.sessionId = checkNotNull(sessionId);
  this.primitiveType = primitiveType;
  this.descriptor = descriptor;
  this.clusterMembershipService = clusterMembershipService;
  this.protocol = protocol;
  this.primaryElection = primaryElection;
  this.threadContext = threadContext;
  clusterMembershipService.addListener(membershipEventListener);
  primaryElection.addListener(primaryElectionListener);
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)
      .addValue(clientName)
      .add("type", primitiveType.name())
      .add("name", descriptor.name())
      .build());
}
 

/**
 * Tests submitting a sequential event.
 */
private void testSequentialEvent(int nodes, int backups, Replication replication) throws Throwable {
  createServers(nodes);

  AtomicLong count = new AtomicLong();
  AtomicLong index = new AtomicLong();

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, backups, replication);
  session.<Long>addEventListener(CHANGE_EVENT, event -> {
    threadAssertEquals(count.incrementAndGet(), 2L);
    threadAssertEquals(index.get(), SERIALIZER.decode(event.value()));
    resume();
  });

  session.execute(operation(EVENT, SERIALIZER.encode(true)))
      .<Long>thenApply(SERIALIZER::decode)
      .thenAccept(result -> {
        threadAssertNotNull(result);
        threadAssertEquals(count.incrementAndGet(), 1L);
        index.set(result);
        resume();
      });

  await(5000, 2);
}
 

/**
 * Tests submitting sequential events to all sessions.
 */
private void testEvents(int nodes, int backups, Replication replication) throws Throwable {
  createServers(nodes);

  PrimaryBackupClient client1 = createClient();
  SessionClient session1 = createProxy(client1, backups, replication);
  session1.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  PrimaryBackupClient client2 = createClient();
  SessionClient session2 = createProxy(client2, backups, replication);
  session2.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  session1.execute(operation(READ, null)).thenRun(this::resume);
  session2.execute(operation(READ, null)).thenRun(this::resume);
  await(5000, 2);

  session1.execute(operation(EVENT, SERIALIZER.encode(false))).thenRun(this::resume);
  await(5000, 3);
}
 

/**
 * Tests submitting a linearizable event that publishes to all sessions.
 */
private void testManyEvents(int nodes, int backups, Replication replication) throws Throwable {
  createServers(nodes);

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, backups, replication);
  session.addEventListener(CHANGE_EVENT, message -> {
    threadAssertNotNull(message);
    resume();
  });

  for (int i = 0; i < 10; i++) {
    session.execute(operation(EVENT, SERIALIZER.encode(true))).thenRun(this::resume);

    await(5000, 2);
  }
}
 

/**
 * Tests events after a primary shutdown.
 */
private void testManyEventsAfterPrimaryShutdown(Replication replication) throws Throwable {
  List<PrimaryBackupServer> servers = createServers(5);

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, 3, replication);
  session.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  for (int i = 0; i < 10; i++) {
    session.execute(operation(EVENT, SERIALIZER.encode(true))).thenRun(this::resume);

    await(5000, 2);
  }

  PrimaryBackupServer leader = servers.stream().filter(s -> s.getRole() == Role.PRIMARY).findFirst().get();
  leader.stop().get(10, TimeUnit.SECONDS);

  for (int i = 0; i < 10; i++) {
    session.execute(operation(EVENT, SERIALIZER.encode(true))).thenRun(this::resume);

    await(5000, 2);
  }
}
 

/**
 * Tests a session closing.
 */
private void testSessionClose(Replication replication) throws Throwable {
  createServers(3);

  PrimaryBackupClient client1 = createClient();
  SessionClient session1 = createProxy(client1, 2, replication);
  PrimaryBackupClient client2 = createClient();
  session1.execute(operation(CLOSE)).thenRun(this::resume);
  await(Duration.ofSeconds(10).toMillis(), 1);
  session1.addEventListener(CLOSE_EVENT, e -> resume());
  SessionClient session2 = createProxy(client2, 2, replication);
  session2.execute(operation(READ)).thenRun(this::resume);
  await(5000);
  session2.close().thenRun(this::resume);
  await(Duration.ofSeconds(10).toMillis(), 2);
}
 

@Override
public <S> ProxyClient<S> newProxy(String primitiveName, PrimitiveType primitiveType, Class<S> serviceType, ServiceConfig serviceConfig, PartitionService partitionService) {
  PartitionGroup partitionGroup = partitionService.getPartitionGroup(this);
  if (partitionGroup == null) {
    throw new ConfigurationException("No Raft partition group matching the configured protocol exists");
  }

  Collection<SessionClient> partitions = partitionGroup.getPartitions().stream()
      .map(partition -> ((RaftPartition) partition).getClient()
          .sessionBuilder(primitiveName, primitiveType, serviceConfig)
          .withMinTimeout(config.getMinTimeout())
          .withMaxTimeout(config.getMaxTimeout())
          .withReadConsistency(config.getReadConsistency())
          .withCommunicationStrategy(config.getCommunicationStrategy())
          .withRecoveryStrategy(config.getRecoveryStrategy())
          .withMaxRetries(config.getMaxRetries())
          .withRetryDelay(config.getRetryDelay())
          .build())
      .collect(Collectors.toList());
  return new DefaultProxyClient<>(primitiveName, primitiveType, this, serviceType, partitions, config.getPartitioner());
}
 

public RecoveringSessionClient(
    String clientId,
    PartitionId partitionId,
    String name,
    PrimitiveType primitiveType,
    Supplier<CompletableFuture<SessionClient>> sessionFactory,
    ThreadContext context) {
  this.partitionId = checkNotNull(partitionId);
  this.name = checkNotNull(name);
  this.primitiveType = checkNotNull(primitiveType);
  this.proxyFactory = checkNotNull(sessionFactory);
  this.context = checkNotNull(context);
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)
      .addValue(clientId)
      .build());
}
 

/**
 * Runs operations for a single Raft proxy.
 */
private void runProxy(SessionClient proxy, CompletableFuture<Void> future) {
  int count = totalOperations.incrementAndGet();
  if (count > TOTAL_OPERATIONS) {
    future.complete(null);
  } else if (count % 10 < WRITE_RATIO) {
    proxy.execute(operation(PUT, CLIENT_SERIALIZER.encode(Maps.immutableEntry(randomKey(), UUID.randomUUID().toString()))))
        .whenComplete((result, error) -> {
          if (error == null) {
            writeCount.incrementAndGet();
          }
          runProxy(proxy, future);
        });
  } else {
    proxy.execute(operation(GET, CLIENT_SERIALIZER.encode(randomKey()))).whenComplete((result, error) -> {
      if (error == null) {
        readCount.incrementAndGet();
      }
      runProxy(proxy, future);
    });
  }
}
 

/**
 * Opens a new client, completing the provided future only once the client has been opened.
 *
 * @param future the future to be completed once the client is opened
 */
private void openProxy(CompletableFuture<SessionClient> future) {
  log.debug("Opening proxy session");
  proxyFactory.get().thenCompose(proxy -> proxy.connect()).whenComplete((proxy, error) -> {
    if (error == null) {
      synchronized (this) {
        this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)
            .addValue(proxy.sessionId())
            .add("type", proxy.type())
            .add("name", proxy.name())
            .build());
        this.session = proxy;
        proxy.addStateChangeListener(this::onStateChange);
        eventListeners.entries().forEach(entry -> proxy.addEventListener(entry.getKey(), entry.getValue()));
        onStateChange(PrimitiveState.CONNECTED);
      }
      future.complete(this);
    } else {
      recoverTask = context.schedule(Duration.ofSeconds(1), () -> openProxy(future));
    }
  });
}
 

@Override
public <S> ProxyClient<S> newProxy(
    String primitiveName,
    PrimitiveType primitiveType,
    Class<S> serviceType,
    ServiceConfig serviceConfig,
    PartitionService partitionService) {
  Collection<SessionClient> partitions = IntStream.range(0, config.getPartitions())
      .mapToObj(partition -> {
        SessionId sessionId = SessionId.from(sessionIds.incrementAndGet());
        return new TestSessionClient(
            primitiveName,
            primitiveType,
            sessionId,
            PartitionId.from(group(), partition),
            new ThreadPoolContext(threadPool),
            registry.getOrCreateService(
                PartitionId.from(group(), partition),
                primitiveName,
                primitiveType,
                serviceConfig));
      })
      .collect(Collectors.toList());
  return new DefaultProxyClient<>(
      primitiveName,
      primitiveType,
      this,
      serviceType,
      partitions,
      config.getPartitioner());
}
 

/**
 * Creates a test session.
 */
private SessionClient createSession(RaftClient client, ReadConsistency consistency) throws Exception {
  return client.sessionBuilder("raft-test", TestPrimitiveType.INSTANCE, new ServiceConfig())
      .withReadConsistency(consistency)
      .withMinTimeout(Duration.ofMillis(250))
      .withMaxTimeout(Duration.ofSeconds(5))
      .build()
      .connect()
      .get(10, TimeUnit.SECONDS);
}
 

/**
 * Tests submitting a command to a partial cluster.
 */
private void testSubmitCommand(int nodes, int backups, Replication replication) throws Throwable {
  createServers(nodes);

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, backups, replication);
  session.execute(operation(WRITE)).thenRun(this::resume);

  await(5000);
}
 

@Override
public synchronized CompletableFuture<SessionClient> connect() {
  CompletableFuture<SessionClient> future = new CompletableFuture<>();
  service.open(sessionId, this).whenCompleteAsync((result, error) -> {
    if (error == null) {
      changeState(PrimitiveState.CONNECTED);
      future.complete(this);
    } else {
      future.completeExceptionally(error);
    }
  }, context);
  return future;
}
 

/**
 * Tests submitting a query with a configured consistency level.
 */
private void testSubmitQuery(int nodes, int backups, Replication replication) throws Throwable {
  createServers(nodes);

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, backups, replication);
  session.execute(operation(READ)).thenRun(this::resume);

  await(5000);
}
 

/**
 * Creates a test session.
 */
private SessionClient createProxy(RaftClient client, ReadConsistency consistency) {
  return client.sessionBuilder("raft-fuzz-test", TestPrimitiveType.INSTANCE, new ServiceConfig())
      .withReadConsistency(consistency)
      .withCommunicationStrategy(COMMUNICATION_STRATEGY)
      .build()
      .connect()
      .join();
}
 

/**
 * Creates a test session.
 */
private SessionClient createProxy(RaftClient client) {
  return client.sessionBuilder("raft-performance-test", TestPrimitiveType.INSTANCE, new ServiceConfig())
      .withReadConsistency(READ_CONSISTENCY)
      .withCommunicationStrategy(COMMUNICATION_STRATEGY)
      .build();
}
 

/**
 * Tests submitting a linearizable event that publishes to all sessions.
 */
private void testManySessionsManyEvents(Replication replication) throws Throwable {
  createServers(3);

  PrimaryBackupClient client = createClient();
  SessionClient session = createProxy(client, 2, replication);
  session.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  SessionClient session1 = createProxy(createClient(), 2, replication);
  session1.execute(operation(READ)).thenRun(this::resume);
  session1.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  SessionClient session2 = createProxy(createClient(), 2, replication);
  session2.execute(operation(READ)).thenRun(this::resume);
  session2.addEventListener(CHANGE_EVENT, event -> {
    threadAssertNotNull(event);
    resume();
  });

  await(5000, 2);

  for (int i = 0; i < 10; i++) {
    session.execute(operation(EVENT, SERIALIZER.encode(false))).thenRun(this::resume);

    await(10000, 4);
  }
}
 

/**
 * Runs a single performance test iteration, returning the iteration run time.
 */
@SuppressWarnings("unchecked")
private long runIteration() throws Exception {
  reset();

  createServers(3);

  CompletableFuture<Void>[] futures = new CompletableFuture[NUM_CLIENTS];
  RaftClient[] clients = new RaftClient[NUM_CLIENTS];
  SessionClient[] proxies = new SessionClient[NUM_CLIENTS];
  for (int i = 0; i < NUM_CLIENTS; i++) {
    CompletableFuture<Void> future = new CompletableFuture<>();
    clients[i] = createClient();
    proxies[i] = createProxy(clients[i]).connect().join();
    futures[i] = future;
  }

  long startTime = System.currentTimeMillis();
  for (int i = 0; i < clients.length; i++) {
    runProxy(proxies[i], futures[i]);
  }
  CompletableFuture.allOf(futures).join();
  long endTime = System.currentTimeMillis();
  long runTime = endTime - startTime;
  System.out.println(String.format("readCount: %d/%d, writeCount: %d/%d, runTime: %dms",
      readCount.get(),
      TOTAL_OPERATIONS,
      writeCount.get(),
      TOTAL_OPERATIONS,
      runTime));
  return runTime;
}
 

/**
 * Creates a new primary-backup proxy.
 */
private SessionClient createProxy(PrimaryBackupClient client, int backups, Replication replication) {
  try {
    return client.sessionBuilder("primary-backup-test", TestPrimitiveType.INSTANCE, new ServiceConfig())
        .withNumBackups(backups)
        .withReplication(replication)
        .build()
        .connect()
        .get(30, TimeUnit.SECONDS);
  } catch (InterruptedException | ExecutionException | TimeoutException e) {
    throw new RuntimeException(e);
  }
}
 

RaftSessionSequencer(RaftSessionState state) {
  this.state = state;
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)
      .addValue(state.getSessionId())
      .add("type", state.getPrimitiveType())
      .add("name", state.getPrimitiveName())
      .build());
}
 

private SessionClient createProxy(RaftClient client) {
  return client.sessionBuilder(ClusterPrimitiveType.PRIMITIVE_NAME,
      ClusterPrimitiveType.INSTANCE, new ServiceConfig())
      .withReadConsistency(ReadConsistency.SEQUENTIAL)
      .withCommunicationStrategy(CommunicationStrategy.LEADER)
      .build()
      .connect()
      .join();
}
 

RaftSessionListener(RaftClientProtocol protocol, MemberSelector memberSelector, RaftSessionState state, RaftSessionSequencer sequencer, Executor executor) {
  this.protocol = checkNotNull(protocol, "protocol cannot be null");
  this.memberSelector = checkNotNull(memberSelector, "nodeSelector cannot be null");
  this.state = checkNotNull(state, "state cannot be null");
  this.sequencer = checkNotNull(sequencer, "sequencer cannot be null");
  this.executor = checkNotNull(executor, "executor cannot be null");
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(SessionClient.class)
      .addValue(state.getSessionId())
      .add("type", state.getPrimitiveType())
      .add("name", state.getPrimitiveName())
      .build());
  protocol.registerPublishListener(state.getSessionId(), this::handlePublish, executor);
}
 

/**
 * Creates a new primitive instance.
 */
private TestPrimitive createPrimitive(RaftClient client, ReadConsistency consistency) throws Exception {
  SessionClient partition = createSession(client, consistency);
  ProxyClient<TestPrimitiveService> proxy = new DefaultProxyClient<>(
      "test",
      TestPrimitiveType.INSTANCE,
      MultiRaftProtocol.builder().build(),
      TestPrimitiveService.class,
      Collections.singletonList(partition),
      (key, partitions) -> partitions.get(0));
  PrimitiveRegistry registry = mock(PrimitiveRegistry.class);
  when(registry.createPrimitive(any(String.class), any(PrimitiveType.class))).thenReturn(CompletableFuture.completedFuture(new PrimitiveInfo("raft-test", TestPrimitiveType.INSTANCE)));
  when(registry.deletePrimitive(any(String.class))).thenReturn(CompletableFuture.completedFuture(null));
  return new TestPrimitiveImpl(proxy, registry);
}
 

/**
 * Tests keeping a client session alive.
 */
@Test
public void testClientKeepAlive() throws Throwable {
  createServers(3);
  RaftClient client = createClient();
  SessionClient session = createSession(client);
  Thread.sleep(Duration.ofSeconds(10).toMillis());
  threadAssertTrue(session.getState() == PrimitiveState.CONNECTED);
}
 

@Override
public CompletableFuture<SessionClient> connect() {
  CompletableFuture<SessionClient> future = new CompletableFuture<>();
  threadContext.execute(() -> {
    connect(1, future);
  });
  return future;
}
 

private SessionClient createProxy(RaftClient client) {
  return client.sessionBuilder(ClusterPrimitiveType.PRIMITIVE_NAME,
      ClusterPrimitiveType.INSTANCE, new ServiceConfig())
      .withReadConsistency(ReadConsistency.SEQUENTIAL)
      .withCommunicationStrategy(CommunicationStrategy.LEADER)
      .build()
      .connect()
      .join();
}