io.atomix.primitive.session.SessionId Java Examples

private RaftSession createSession(long sessionId) {
  RaftServiceContext context = mock(RaftServiceContext.class);
  when(context.serviceType()).thenReturn(TestPrimitiveType.instance());
  when(context.serviceName()).thenReturn("test");
  when(context.serviceId()).thenReturn(PrimitiveId.from(1));

  RaftContext server = mock(RaftContext.class);
  when(server.getProtocol()).thenReturn(mock(RaftServerProtocol.class));
  RaftServiceManager manager = mock(RaftServiceManager.class);
  when(manager.executor()).thenReturn(mock(ThreadContext.class));
  when(server.getServiceManager()).thenReturn(manager);

  return new RaftSession(
      SessionId.from(sessionId),
      MemberId.from("1"),
      "test",
      TestPrimitiveType.instance(),
      ReadConsistency.LINEARIZABLE,
      100,
      5000,
      System.currentTimeMillis(),
      Serializer.using(Namespaces.BASIC),
      context,
      server,
      mock(ThreadContextFactory.class));
}
 

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventAtCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(2)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  CommandResponse response = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(2)
      .build();

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(1, run.getAndIncrement()));
  sequencer.sequenceEvent(request, () -> assertEquals(0, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

public AbstractAtomicMapService(PrimitiveType primitiveType) {
  super(primitiveType, AtomicMapClient.class);
  serializer = Serializer.using(Namespace.builder()
      .register(primitiveType.namespace())
      .register(SessionId.class)
      .register(TransactionId.class)
      .register(TransactionScope.class)
      .register(MapEntryValue.class)
      .register(MapEntryValue.Type.class)
      .register(new HashMap().keySet().getClass())
      .register(DefaultIterator.class)
      .register(LockContext.class)
      .register(LockHolder.class)
      .build());
  map = createMap();
}
 

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 query to the cluster.
 */
@Test
public void testSubmitQuery() throws Throwable {
  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  when(connection.query(any(QueryRequest.class)))
      .thenReturn(CompletableFuture.completedFuture(QueryResponse.builder()
          .withStatus(RaftResponse.Status.OK)
          .withIndex(10)
          .withResult("Hello world!".getBytes())
          .build()));

  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  RaftSessionManager manager = mock(RaftSessionManager.class);
  ThreadContext threadContext = new TestContext();

  RaftSessionInvoker submitter = new RaftSessionInvoker(mock(RaftSessionConnection.class), connection, state, new RaftSessionSequencer(state), manager, threadContext);
  assertTrue(Arrays.equals(submitter.invoke(operation(QUERY)).get(), "Hello world!".getBytes()));
  assertEquals(10, state.getResponseIndex());
}
 

public DistributedLogSession(
    PartitionId partitionId,
    SessionId sessionId,
    ClusterMembershipService clusterMembershipService,
    LogClientProtocol protocol,
    PrimaryElection primaryElection,
    ThreadContext threadContext) {
  this.partitionId = checkNotNull(partitionId, "partitionId cannot be null");
  this.sessionId = checkNotNull(sessionId, "sessionId cannot be null");
  this.protocol = checkNotNull(protocol, "protocol cannot be null");
  this.primaryElection = checkNotNull(primaryElection, "primaryElection cannot be null");
  this.threadContext = checkNotNull(threadContext, "threadContext cannot be null");
  this.memberId = clusterMembershipService.getLocalMember().id();
  this.subject = String.format("%s-%s-%s", partitionId.group(), partitionId.id(), sessionId);
  clusterMembershipService.addListener(membershipEventListener);
  primaryElection.addListener(primaryElectionListener);
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(DistributedLogProducer.class)
      .addValue(partitionId.group() != null
          ? String.format("%s-%d", partitionId.group(), partitionId.id())
          : partitionId.id())
      .build());
}
 

@Test
@SuppressWarnings("unchecked")
public void testSnapshot() throws Exception {
  ServiceContext context = mock(ServiceContext.class);
  when(context.serviceType()).thenReturn(DistributedSetType.instance());
  when(context.serviceName()).thenReturn("test");
  when(context.serviceId()).thenReturn(PrimitiveId.from(1));
  when(context.wallClock()).thenReturn(new WallClock());

  Session session = mock(Session.class);
  when(session.sessionId()).thenReturn(SessionId.from(1));

  DefaultDistributedSetService service = new DefaultDistributedSetService();
  service.init(context);

  service.add("foo");

  Buffer buffer = HeapBuffer.allocate();
  service.backup(new DefaultBackupOutput(buffer, service.serializer()));

  service = new DefaultDistributedSetService();
  service.restore(new DefaultBackupInput(buffer.flip(), service.serializer()));

  assertTrue(service.contains("foo"));
}
 

@Test
@SuppressWarnings("unchecked")
public void testSnapshot() throws Exception {
  ServiceContext context = mock(ServiceContext.class);
  when(context.serviceType()).thenReturn(DistributedSetType.instance());
  when(context.serviceName()).thenReturn("test");
  when(context.serviceId()).thenReturn(PrimitiveId.from(1));
  when(context.wallClock()).thenReturn(new WallClock());

  Session session = mock(Session.class);
  when(session.sessionId()).thenReturn(SessionId.from(1));

  DefaultDistributedListService service = new DefaultDistributedListService();
  service.init(context);

  service.add("foo");

  Buffer buffer = HeapBuffer.allocate();
  service.backup(new DefaultBackupOutput(buffer, service.serializer()));

  service = new DefaultDistributedListService();
  service.restore(new DefaultBackupInput(buffer.flip(), service.serializer()));

  assertEquals("foo", service.get(0));
}
 

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventAbsentCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));

  PublishRequest request1 = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(2)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  PublishRequest request2 = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(3)
      .withPreviousIndex(2)
      .withEvents(Collections.emptyList())
      .build();

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceEvent(request1, () -> assertEquals(0, run.getAndIncrement()));
  sequencer.sequenceEvent(request2, () -> assertEquals(1, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

public DistributedLogSessionClient(
    String clientName,
    PartitionId partitionId,
    ClusterMembershipService clusterMembershipService,
    LogClientProtocol protocol,
    Supplier<CompletableFuture<SessionId>> sessionIdProvider,
    PrimaryElection primaryElection,
    ThreadContextFactory threadContextFactory,
    boolean closeOnStop) {
  this.clientName = clientName;
  this.partitionId = partitionId;
  this.clusterMembershipService = clusterMembershipService;
  this.protocol = protocol;
  this.sessionIdProvider = sessionIdProvider;
  this.primaryElection = primaryElection;
  this.threadContextFactory = threadContextFactory;
  this.threadContext = threadContextFactory.createContext();
  this.closeOnStop = closeOnStop;
}
 

/**
 * Tests sequencing callbacks with the sequencer.
 */
@Test
public void testSequenceResponses() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence1 = sequencer.nextRequest();
  long sequence2 = sequencer.nextRequest();
  assertTrue(sequence2 == sequence1 + 1);

  CommandResponse commandResponse = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(0)
      .build();

  QueryResponse queryResponse = QueryResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(0)
      .build();

  AtomicBoolean run = new AtomicBoolean();
  sequencer.sequenceResponse(sequence2, queryResponse, () -> run.set(true));
  sequencer.sequenceResponse(sequence1, commandResponse, () -> assertFalse(run.get()));
  assertTrue(run.get());
}
 

/**
 * Tests updating client session state.
 */
@Test
public void testSessionState() {
  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  assertEquals(state.getSessionId(), SessionId.from(1));
  assertEquals(1, state.getResponseIndex());
  assertEquals(1, state.getEventIndex());
  state.setCommandRequest(2);
  assertEquals(2, state.getCommandRequest());
  assertEquals(3, state.nextCommandRequest());
  assertEquals(3, state.getCommandRequest());
  state.setCommandResponse(3);
  assertEquals(3, state.getCommandResponse());
  state.setResponseIndex(4);
  assertEquals(4, state.getResponseIndex());
  state.setResponseIndex(3);
  assertEquals(4, state.getResponseIndex());
  state.setEventIndex(5);
  assertEquals(5, state.getEventIndex());
}
 

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventBeforeCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(1)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  CommandResponse response = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(1)
      .build();

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceEvent(request, () -> assertEquals(0, run.getAndIncrement()));
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(1, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

/**
 * Resets indexes for the given session.
 *
 * @param sessionId The session for which to reset indexes.
 * @return A completable future to be completed once the session's indexes have been reset.
 */
CompletableFuture<Void> resetIndexes(SessionId sessionId) {
  RaftSessionState sessionState = sessions.get(sessionId.id());
  if (sessionState == null) {
    return Futures.exceptionalFuture(new IllegalArgumentException("Unknown session: " + sessionId));
  }

  CompletableFuture<Void> future = new CompletableFuture<>();

  KeepAliveRequest request = KeepAliveRequest.builder()
      .withSessionIds(new long[]{sessionId.id()})
      .withCommandSequences(new long[]{sessionState.getCommandResponse()})
      .withEventIndexes(new long[]{sessionState.getEventIndex()})
      .build();

  connection.keepAlive(request).whenComplete((response, error) -> {
    if (error == null) {
      if (response.status() == RaftResponse.Status.OK) {
        future.complete(null);
      } else {
        future.completeExceptionally(response.error().createException());
      }
    } else {
      future.completeExceptionally(error);
    }
  });
  return future;
}
 

@Test
public void testSnapshot() throws Exception {
  ServiceContext context = mock(ServiceContext.class);
  when(context.serviceType()).thenReturn(WorkQueueType.instance());
  when(context.serviceName()).thenReturn("test");
  when(context.serviceId()).thenReturn(PrimitiveId.from(1));

  Session session = mock(Session.class);
  when(session.sessionId()).thenReturn(SessionId.from(1));
  when(context.currentSession()).thenReturn(session);

  DefaultWorkQueueService service = new DefaultWorkQueueService();
  service.init(context);
  service.register(session);

  service.add(Arrays.asList("Hello world!".getBytes()));

  Buffer buffer = HeapBuffer.allocate();
  service.backup(new DefaultBackupOutput(buffer, service.serializer()));

  service = new DefaultWorkQueueService();
  service.init(context);
  service.register(session);
  service.restore(new DefaultBackupInput(buffer.flip(), service.serializer()));

  Collection<Task<byte[]>> value = service.take(1);
  assertNotNull(value);
  assertEquals(1, value.size());
  assertArrayEquals("Hello world!".getBytes(), value.iterator().next().payload());
}
 

/**
 * Creates a service session.
 *
 * @param sessionId the session to create
 * @param memberId  the owning node ID
 * @return the service session
 */
public PrimaryBackupSession createSession(long sessionId, MemberId memberId) {
  PrimaryBackupSession session = new PrimaryBackupSession(SessionId.from(sessionId), memberId, service.serializer(), this);
  if (sessions.putIfAbsent(sessionId, session) == null) {
    service.register(session);
  }
  return session;
}
 

@Override
public void registerResetListener(SessionId sessionId,
                                  Consumer<ResetRequest> listener, Executor executor) {
  messagingService.registerHandler(String.format("reset-%d", sessionId.id()), (e, p) -> {
    listener.accept(serializer.decode(p));
  }, executor);
}
 

@Test
public void testAddRemoveSession() throws Exception {
  RaftSessionRegistry sessionManager = new RaftSessionRegistry();
  RaftSession session = createSession(1);
  sessionManager.addSession(session);
  assertNotNull(sessionManager.getSession(1));
  assertNotNull(sessionManager.getSession(session.sessionId()));
  assertEquals(0, sessionManager.getSessions(PrimitiveId.from(1)).size());
  session.open();
  assertEquals(1, sessionManager.getSessions(PrimitiveId.from(1)).size());
  sessionManager.removeSession(SessionId.from(1));
  assertNull(sessionManager.getSession(1));
}
 

@Override
public final void close(SessionId sessionId) {
  Session session = sessions.remove(sessionId);
  if (session != null) {
    onClose(session);
  }
}
 

public DefaultDistributedMultisetService() {
  super(DistributedMultisetType.instance(), HashMultiset.create());
  this.serializer = Serializer.using(Namespace.builder()
      .register(DistributedMultisetType.instance().namespace())
      .register(SessionId.class)
      .register(DefaultDistributedCollectionService.IteratorContext.class)
      .register(IteratorContext.class)
      .build());
}
 

@Override
public void registerPublishListener(SessionId sessionId, Consumer<PublishRequest> listener,
                                    Executor executor) {
  messagingService.registerHandler(String.format("publish-%d", sessionId.id()), (e, p) -> {
    listener.accept(serializer.decode(p));
  }, executor);
}
 

public PrimaryBackupSession(SessionId sessionId, MemberId memberId, Serializer serializer, PrimaryBackupServiceContext context) {
  super(sessionId, context.serviceName(), context.serviceType(), memberId, serializer);
  this.context = context;
  this.log = ContextualLoggerFactory.getLogger(getClass(), LoggerContext.builder(getClass())
      .addValue(context.serverName())
      .add("session", sessionId)
      .build());
}
 

/**
 * Closes a session.
 *
 * @param sessionId the session identifier.
 * @param delete whether to delete the service
 * @return A completable future to be completed once the session is closed.
 */
public CompletableFuture<Void> closeSession(SessionId sessionId, boolean delete) {
  RaftSessionState state = sessions.get(sessionId.id());
  if (state == null) {
    return Futures.exceptionalFuture(new RaftException.UnknownSession("Unknown session: " + sessionId));
  }

  log.debug("Closing session {}", sessionId);
  CloseSessionRequest request = CloseSessionRequest.builder()
      .withSession(sessionId.id())
      .withDelete(delete)
      .build();

  CompletableFuture<Void> future = new CompletableFuture<>();
  connection.closeSession(request).whenComplete((response, error) -> {
    sessions.remove(sessionId.id());
    if (error == null) {
      if (response.status() == RaftResponse.Status.OK) {
        future.complete(null);
      } else {
        future.completeExceptionally(response.error().createException());
      }
    } else {
      future.completeExceptionally(error);
    }
  });
  return future;
}
 

@Override
public void registerPublishListener(SessionId sessionId, Consumer<PublishRequest> listener,
                                    Executor executor) {
  messagingService.registerHandler(String.format("publish-%d", sessionId.id()), (e, p) -> {
    listener.accept(serializer.decode(p));
  }, executor);
}
 

TestSessionClient(
    String name,
    PrimitiveType type,
    SessionId sessionId,
    PartitionId partitionId,
    ThreadContext context,
    TestProtocolService service) {
  this.name = name;
  this.type = type;
  this.sessionId = sessionId;
  this.partitionId = partitionId;
  this.context = context;
  this.service = service;
}
 

@Test
public void testSnapshot() throws Exception {
  ServiceContext context = mock(ServiceContext.class);
  when(context.serviceType()).thenReturn(LeaderElectionType.instance());
  when(context.serviceName()).thenReturn("test");
  when(context.serviceId()).thenReturn(PrimitiveId.from(1));
  when(context.wallClock()).thenReturn(new WallClock());

  Session session = mock(Session.class);
  when(session.sessionId()).thenReturn(SessionId.from(1));
  when(context.currentSession()).thenReturn(session);

  DefaultLeaderElectionService service = new DefaultLeaderElectionService();
  service.init(context);
  service.register(session);

  byte[] id = "a".getBytes();
  service.run(id);

  Buffer buffer = HeapBuffer.allocate();
  service.backup(new DefaultBackupOutput(buffer, service.serializer()));

  service = new DefaultLeaderElectionService();
  service.init(context);
  service.register(session);
  service.restore(new DefaultBackupInput(buffer.flip(), service.serializer()));

  Leadership<byte[]> value = service.getLeadership();
  assertNotNull(value);
  assertArrayEquals(id, value.leader().id());
}
 

/**
 * Expires the given session.
 *
 * @param sessionId the session identifier
 * @return a future to be completed once the session has been expired
 */
public CompletableFuture<Void> expire(SessionId sessionId) {
  CompletableFuture<Void> future = new CompletableFuture<>();
  context.execute(() -> {
    TestProtocolSession session = sessions.remove(sessionId);
    if (session != null) {
      session.setState(Session.State.EXPIRED);
      service.expire(sessionId);
      future.complete(null);
    } else {
      future.completeExceptionally(new PrimitiveException.UnknownSession());
    }
  });
  return future;
}
 

Waiter(SessionId session, long index, long id, int acquirePermits, long expire) {
  this.session = session;
  this.index = index;
  this.id = id;
  this.acquirePermits = acquirePermits;
  this.expire = expire;
}
 

/**
 * Closes the given session.
 *
 * @param sessionId the session identifier
 * @return a future to be completed once the session has been closed
 */
public CompletableFuture<Void> close(SessionId sessionId) {
  CompletableFuture<Void> future = new CompletableFuture<>();
  context.execute(() -> {
    TestProtocolSession session = sessions.remove(sessionId);
    if (session != null) {
      session.setState(Session.State.CLOSED);
      service.close(sessionId);
      future.complete(null);
    } else {
      future.completeExceptionally(new PrimitiveException.UnknownSession());
    }
  });
  return future;
}
 

/**
 * Executes the given operation.
 *
 * @param sessionId the session performing the operation
 * @param operation the operation to execute
 * @return a future to be completed with the operation result
 */
public CompletableFuture<byte[]> execute(SessionId sessionId, PrimitiveOperation operation) {
  CompletableFuture<byte[]> future = new CompletableFuture<>();
  context.execute(() -> {
    Session session = sessions.get(sessionId);
    if (session == null) {
      future.completeExceptionally(new PrimitiveException.UnknownSession());
    }

    long timestamp = timestamp();
    this.session = session;
    this.operationType = operation.id().type();
    service.tick(new WallClockTimestamp(timestamp));

    try {
      byte[] result = service.apply(new DefaultCommit<>(
          index.incrementAndGet(),
          operation.id(),
          operation.value(),
          session,
          timestamp));
      future.complete(result);
    } catch (Exception e) {
      future.completeExceptionally(new PrimitiveException.ServiceException(e.getMessage()));
    }
  });
  return future;
}