io.atomix.primitive.partition.PartitionId Java Examples

public HashBasedPrimaryElection(
    PartitionId partitionId,
    ClusterMembershipService clusterMembershipService,
    PartitionGroupMembershipService groupMembershipService,
    ClusterCommunicationService communicationService,
    ScheduledExecutorService executor) {
  this.partitionId = partitionId;
  this.clusterMembershipService = clusterMembershipService;
  this.groupMembershipService = groupMembershipService;
  this.communicationService = communicationService;
  this.subject = String.format("primary-election-counter-%s-%d", partitionId.group(), partitionId.id());
  recomputeTerm(groupMembershipService.getMembership(partitionId.group()));
  groupMembershipService.addListener(groupMembershipEventListener);
  clusterMembershipService.addListener(clusterMembershipEventListener);
  communicationService.subscribe(subject, SERIALIZER::decode, this::updateCounters, executor);
  broadcastFuture = executor.scheduleAtFixedRate(this::broadcastCounters, BROADCAST_INTERVAL, BROADCAST_INTERVAL, TimeUnit.MILLISECONDS);
}
 

public DefaultRaftSessionClient(
    String serviceName,
    PrimitiveType primitiveType,
    ServiceConfig serviceConfig,
    PartitionId partitionId,
    RaftClientProtocol protocol,
    MemberSelectorManager selectorManager,
    RaftSessionManager sessionManager,
    ReadConsistency readConsistency,
    CommunicationStrategy communicationStrategy,
    ThreadContext context,
    Duration minTimeout,
    Duration maxTimeout) {
  this.serviceName = checkNotNull(serviceName, "serviceName cannot be null");
  this.primitiveType = checkNotNull(primitiveType, "serviceType cannot be null");
  this.serviceConfig = checkNotNull(serviceConfig, "serviceConfig cannot be null");
  this.partitionId = checkNotNull(partitionId, "partitionId cannot be null");
  this.protocol = checkNotNull(protocol, "protocol cannot be null");
  this.selectorManager = checkNotNull(selectorManager, "selectorManager cannot be null");
  this.readConsistency = checkNotNull(readConsistency, "readConsistency cannot be null");
  this.communicationStrategy = checkNotNull(communicationStrategy, "communicationStrategy cannot be null");
  this.context = checkNotNull(context, "context cannot be null");
  this.minTimeout = checkNotNull(minTimeout, "minTimeout cannot be null");
  this.maxTimeout = checkNotNull(maxTimeout, "maxTimeout cannot be null");
  this.sessionManager = checkNotNull(sessionManager, "sessionManager cannot be null");
}
 

/**
 * Handles a partition proxy state change.
 */
private synchronized void onStateChange(PartitionId partitionId, PrimitiveState state) {
  states.put(partitionId, state);
  switch (state) {
    case CONNECTED:
      if (this.state != PrimitiveState.CONNECTED && !states.containsValue(PrimitiveState.SUSPENDED) && !states.containsValue(PrimitiveState.CLOSED)) {
        this.state = PrimitiveState.CONNECTED;
        stateChangeListeners.forEach(l -> l.accept(PrimitiveState.CONNECTED));
      }
      break;
    case SUSPENDED:
      if (this.state == PrimitiveState.CONNECTED) {
        this.state = PrimitiveState.SUSPENDED;
        stateChangeListeners.forEach(l -> l.accept(PrimitiveState.SUSPENDED));
      }
      break;
    case CLOSED:
      if (this.state != PrimitiveState.CLOSED) {
        this.state = PrimitiveState.CLOSED;
        stateChangeListeners.forEach(l -> l.accept(PrimitiveState.CLOSED));
      }
      break;
  }
}
 

@Override
public CompletableFuture<Boolean> prepare(TransactionLog<MapUpdate<K, byte[]>> transactionLog) {
  Map<PartitionId, List<MapUpdate<K, byte[]>>> updatesGroupedByMap = Maps.newIdentityHashMap();
  transactionLog.records().forEach(update -> {
    updatesGroupedByMap.computeIfAbsent(getPartition(update.key()), k -> Lists.newLinkedList()).add(update);
  });
  Map<PartitionId, TransactionLog<MapUpdate<K, byte[]>>> transactionsByMap =
      Maps.transformValues(updatesGroupedByMap, list -> new TransactionLog<>(transactionLog.transactionId(), transactionLog.version(), list));

  return Futures.allOf(transactionsByMap.entrySet()
      .stream()
      .map(e -> getProxyClient().applyOn(e.getKey(), service -> service.prepare(e.getValue()))
          .thenApply(v -> v == PrepareResult.OK || v == PrepareResult.PARTIAL_FAILURE))
      .collect(Collectors.toList()))
      .thenApply(list -> list.stream().reduce(Boolean::logicalAnd).orElse(true));
}
 

private synchronized void onStateChange(PartitionId partitionId, PrimitiveState state) {
  states.put(partitionId, state);
  switch (state) {
    case CONNECTED:
      if (!states.containsValue(PrimitiveState.SUSPENDED) && !states.containsValue(PrimitiveState.CLOSED)) {
        changeState(PrimitiveState.CONNECTED);
      }
      break;
    case SUSPENDED:
      changeState(PrimitiveState.SUSPENDED);
      break;
    case CLOSED:
      changeState(PrimitiveState.CLOSED);
      break;
  }
}
 

public DefaultRaftClient(
    String clientId,
    PartitionId partitionId,
    MemberId memberId,
    Collection<MemberId> cluster,
    RaftClientProtocol protocol,
    ThreadContextFactory threadContextFactory,
    boolean closeThreadFactoryOnClose) {
  this.clientId = checkNotNull(clientId, "clientId cannot be null");
  this.partitionId = checkNotNull(partitionId, "partitionId cannot be null");
  this.cluster = checkNotNull(cluster, "cluster cannot be null");
  this.protocol = checkNotNull(protocol, "protocol cannot be null");
  this.threadContextFactory = checkNotNull(threadContextFactory, "threadContextFactory cannot be null");
  this.threadContext = threadContextFactory.createContext();
  this.metadata = new DefaultRaftMetadataClient(clientId, protocol, selectorManager, threadContextFactory.createContext());
  this.sessionManager = new RaftSessionManager(clientId, memberId, protocol, selectorManager, threadContextFactory);
  this.closeThreadFactoryOnClose = closeThreadFactoryOnClose;
}
 

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;
}
 

public PrimaryBackupClient(
    String clientName,
    PartitionId partitionId,
    ClusterMembershipService clusterMembershipService,
    PrimaryBackupClientProtocol protocol,
    PrimaryElection primaryElection,
    SessionIdService sessionIdService,
    ThreadContextFactory threadContextFactory,
    boolean closeOnStop) {
  this.clientName = clientName;
  this.partitionId = partitionId;
  this.clusterMembershipService = clusterMembershipService;
  this.protocol = protocol;
  this.primaryElection = primaryElection;
  this.sessionIdService = sessionIdService;
  this.threadContextFactory = threadContextFactory;
  this.threadContext = threadContextFactory.createContext();
  this.closeOnStop = closeOnStop;
}
 

TestProtocolService(
    PartitionId partition,
    String name,
    PrimitiveType primitiveType,
    ServiceConfig config,
    PrimitiveService service,
    TestProtocolServiceRegistry registry,
    ThreadContext context) {
  this.partition = partition;
  this.name = name;
  this.primitiveType = primitiveType;
  this.config = config;
  this.service = service;
  this.registry = registry;
  this.context = context;
  this.clock = context.schedule(Duration.ofMillis(100), Duration.ofMillis(100), this::tick);
  open();
}
 

@Override
public CompletableFuture<Boolean> prepare(TransactionLog<SetUpdate<String>> transactionLog) {
  Map<PartitionId, List<SetUpdate<String>>> updatesGroupedBySet = Maps.newIdentityHashMap();
  transactionLog.records().forEach(update -> {
    updatesGroupedBySet.computeIfAbsent(getProxyClient().getPartitionId(update.element()), k -> Lists.newLinkedList()).add(update);
  });
  Map<PartitionId, TransactionLog<SetUpdate<String>>> transactionsBySet =
      Maps.transformValues(updatesGroupedBySet, list -> new TransactionLog<>(transactionLog.transactionId(), transactionLog.version(), list));

  return Futures.allOf(transactionsBySet.entrySet()
      .stream()
      .map(e -> getProxyClient().applyOn(e.getKey(), service -> service.prepare(e.getValue()))
          .thenApply(v -> v == PrepareResult.OK || v == PrepareResult.PARTIAL_FAILURE))
      .collect(Collectors.toList()))
      .thenApply(list -> list.stream().reduce(Boolean::logicalAnd).orElse(true));
}
 

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());
}
 

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());
}
 

public RaftPartition(
    PartitionId partitionId,
    RaftPartitionGroupConfig config,
    File dataDirectory,
    ThreadContextFactory threadContextFactory) {
  this.partitionId = partitionId;
  this.config = config;
  this.dataDirectory = dataDirectory;
  this.threadContextFactory = threadContextFactory;
}
 

@Override
public CompletableFuture<Boolean> containsAll(Collection<? extends String> c) {
  Map<PartitionId, Collection<String>> partitions = Maps.newHashMap();
  c.forEach(key -> partitions.computeIfAbsent(getProxyClient().getPartitionId(key), k -> Lists.newArrayList()).add(key));
  return Futures.allOf(partitions.entrySet().stream()
      .map(entry -> getProxyClient()
          .applyOn(entry.getKey(), service -> service.containsAll(entry.getValue())))
      .collect(Collectors.toList()))
      .thenApply(results -> results.stream().reduce(Boolean::logicalAnd).orElse(true));
}
 

private static Collection<RaftPartition> buildPartitions(
    RaftPartitionGroupConfig config,
    ThreadContextFactory threadContextFactory) {
  File partitionsDir = new File(config.getStorageConfig().getDirectory(config.getName()), "partitions");
  List<RaftPartition> partitions = new ArrayList<>(config.getPartitions());
  for (int i = 0; i < config.getPartitions(); i++) {
    partitions.add(new RaftPartition(
        PartitionId.from(config.getName(), i + 1),
        config,
        new File(partitionsDir, String.valueOf(i + 1)),
        threadContextFactory));
  }
  return partitions;
}
 

@Override
public CompletableFuture<Boolean> addAll(Collection<? extends String> c) {
  Map<PartitionId, Collection<String>> partitions = Maps.newHashMap();
  c.forEach(key -> partitions.computeIfAbsent(getProxyClient().getPartitionId(key), k -> Lists.newArrayList()).add(key));
  return Futures.allOf(partitions.entrySet().stream()
      .map(entry -> getProxyClient()
          .applyOn(entry.getKey(), service -> service.addAll(entry.getValue()))
          .thenCompose(result -> checkLocked(result)))
      .collect(Collectors.toList()))
      .thenApply(results -> results.stream().reduce(Boolean::logicalOr).orElse(false));
}
 

private static Collection<LogPartition> buildPartitions(LogPartitionGroupConfig config) {
  List<LogPartition> partitions = new ArrayList<>(config.getPartitions());
  for (int i = 0; i < config.getPartitions(); i++) {
    partitions.add(new LogPartition(PartitionId.from(config.getName(), i + 1), config));
  }
  return partitions;
}
 

/**
 * Removes the given service.
 *
 * @param partitionId the partition identifier
 * @param name the service name
 */
public void removeService(PartitionId partitionId, String name) {
  Map<String, TestProtocolService> services = partitions.get(partitionId);
  if (services != null) {
    services.remove(name);
  }
}
 

@Override
public CompletableFuture<Boolean> containsAll(Collection<? extends String> keys) {
  Map<PartitionId, Collection<String>> partitions = Maps.newHashMap();
  keys.forEach(key -> partitions.computeIfAbsent(getProxyClient().getPartitionId(key), k -> Lists.newArrayList()).add(key));
  return Futures.allOf(partitions.entrySet().stream()
      .map(entry -> getProxyClient()
          .applyOn(entry.getKey(), service -> service.containsKeys(entry.getValue())))
      .collect(Collectors.toList()))
      .thenApply(results -> results.stream().reduce(Boolean::logicalAnd).orElse(false));
}
 

private RaftClient createClient(List<RaftMember> members) throws Throwable {
  final MemberId memberId = nextNodeId();
  final List<MemberId> memberIds = members.stream().map(RaftMember::memberId).collect(Collectors.toList());
  final RaftClient client = RaftClient.builder()
          .withMemberId(memberId)
          .withPartitionId(PartitionId.from("test", 1))
          .withProtocol(protocolFactory.newClientProtocol(memberId))
          .build();
  client.connect(memberIds).thenRun(this::resume);
  await(30000);
  clients.add(client);
  return client;
}
 

public ProxyIterator(
    ProxyClient<S> client,
    PartitionId partitionId,
    OpenFunction<S, T> openFunction,
    NextFunction<S, T> nextFunction,
    CloseFunction<S> closeFunction) {
  this.client = client;
  this.partitionId = partitionId;
  this.nextFunction = nextFunction;
  this.closeFunction = closeFunction;
  this.openFuture = OrderedFuture.wrap(client.applyOn(partitionId, openFunction::open));
  this.batch = openFuture;
}
 

ElectionState(
    PartitionId partitionId,
    List<Registration> registrations,
    Registration primary,
    long term,
    long termStartTime,
    Map<PartitionId, ElectionState> elections) {
  this.partitionId = partitionId;
  this.registrations = Lists.newArrayList(registrations);
  this.primary = primary;
  this.term = term;
  this.termStartTime = termStartTime;
  this.elections = elections;
}
 

@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());
}
 

ElectionState(
    PartitionId partitionId,
    Registration registration,
    Map<PartitionId, ElectionState> elections) {
  registrations = Arrays.asList(registration);
  termStartTime = System.currentTimeMillis();
  primary = registration;
  this.partitionId = partitionId;
  this.term = 1;
  this.elections = elections;
}
 

/**
 * Creates a Raft client.
 */
private PrimaryBackupClient createClient() throws Throwable {
  MemberId memberId = nextMemberId();
  PrimaryBackupClient client = PrimaryBackupClient.builder()
      .withClientName("test")
      .withPartitionId(PartitionId.from("test", 1))
      .withMembershipService(new TestClusterMembershipService(memberId, nodes))
      .withSessionIdProvider(() -> CompletableFuture.completedFuture(nextSessionId()))
      .withPrimaryElection(election)
      .withProtocol(protocolFactory.newClientProtocol(memberId))
      .build();
  clients.add(client);
  return client;
}
 

private void onSessionEnd(Session session) {
  listeners.remove(session.sessionId().id());
  Set<PartitionId> partitions = elections.keySet();
  partitions.forEach(partitionId -> {
    PrimaryTerm oldTerm = term(partitionId);
    elections.compute(partitionId, (k, v) -> v.cleanup(session));
    PrimaryTerm newTerm = term(partitionId);
    if (!Objects.equals(oldTerm, newTerm)) {
      notifyTermChange(partitionId, newTerm);
      scheduleRebalance();
    }
  });
}
 

private static Collection<PrimaryBackupPartition> buildPartitions(PrimaryBackupPartitionGroupConfig config) {
  List<PrimaryBackupPartition> partitions = new ArrayList<>(config.getPartitions());
  for (int i = 0; i < config.getPartitions(); i++) {
    partitions.add(new PrimaryBackupPartition(PartitionId.from(config.getName(), i + 1), config.getMemberGroupProvider()));
  }
  return partitions;
}
 

/**
 * Applies an {@link PrimaryElectorOperations.GetTerm} commit.
 *
 * @param commit GetLeadership commit
 * @return leader
 */
protected PrimaryTerm getTerm(Commit<? extends PrimaryElectorOperations.GetTerm> commit) {
  PartitionId partitionId = commit.value().partitionId();
  try {
    return term(partitionId);
  } catch (Exception e) {
    getLogger().error("State machine operation failed", e);
    throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

/**
 * Applies an {@link PrimaryElectorOperations.Enter} commit.
 *
 * @param commit commit entry
 * @return topic leader. If no previous leader existed this is the node that just entered the race.
 */
protected PrimaryTerm enter(Commit<? extends PrimaryElectorOperations.Enter> commit) {
  try {
    PartitionId partitionId = commit.value().partitionId();
    PrimaryTerm oldTerm = term(partitionId);
    Registration registration = new Registration(
        commit.value().member(),
        commit.session().sessionId().id());
    PrimaryTerm newTerm = elections.compute(partitionId, (k, v) -> {
      if (v == null) {
        return new ElectionState(partitionId, registration, elections);
      } else {
        if (!v.isDuplicate(registration)) {
          return new ElectionState(v).addRegistration(registration);
        } else {
          return v;
        }
      }
    })
        .term();

    if (!Objects.equals(oldTerm, newTerm)) {
      notifyTermChange(partitionId, newTerm);
      scheduleRebalance();
    }
    return newTerm;
  } catch (Exception e) {
    getLogger().error("State machine operation failed", e);
    throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

public DefaultPrimaryElection(PartitionId partitionId, SessionClient proxy, PrimaryElectionService service) {
  this.partitionId = checkNotNull(partitionId);
  this.proxy = proxy;
  this.service = service;
  this.eventListener = event -> {
    if (event.partitionId().equals(partitionId)) {
      listeners.forEach(l -> l.event(event));
    }
  };
  service.addListener(eventListener);
}