io.activej.crdt.storage.cluster.RendezvousPartitionScheme Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of activej-crdt Show documentation
Conflict-free replicated data type implementation
The newest version!
package io.activej.crdt.storage.cluster;

import io.activej.common.builder.AbstractBuilder;
import io.activej.crdt.storage.ICrdtStorage;
import io.activej.crdt.storage.cluster.IDiscoveryService.PartitionScheme;
import io.activej.rpc.client.sender.strategy.RpcStrategies;
import io.activej.rpc.client.sender.strategy.RpcStrategy;
import io.activej.rpc.client.sender.strategy.impl.RendezvousHashing;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.annotations.VisibleForTesting;

import java.util.*;
import java.util.function.Function;
import java.util.function.ToIntFunction;

import static io.activej.common.collection.CollectionUtils.difference;
import static io.activej.crdt.storage.cluster.RendezvousHashSharder.NUMBER_OF_BUCKETS;
import static java.util.stream.Collectors.toSet;

public final class RendezvousPartitionScheme implements PartitionScheme
 {
	private final List> partitionGroups = new ArrayList<>();
	private ToIntFunction keyHashFn = Object::hashCode;
	@SuppressWarnings("unchecked")
	private Function partitionIdGetter = (Function) Function.identity();
	private Function rpcProvider;
	private Function> crdtProvider;

	@SafeVarargs
	public static 
 RendezvousPartitionScheme
 create(RendezvousPartitionGroup
... partitionGroups) {
		return builder(partitionGroups).build();
	}

	public static 
 RendezvousPartitionScheme
 create(List> partitionGroups) {
		return builder(partitionGroups).build();
	}

	@SafeVarargs
	public static 
 RendezvousPartitionScheme
.Builder builder(RendezvousPartitionGroup
... partitionGroups) {
		return builder(List.of(partitionGroups));
	}

	public static 
 RendezvousPartitionScheme
.Builder builder(List> partitionGroups) {
		RendezvousPartitionScheme
 scheme = new RendezvousPartitionScheme<>();
		scheme.partitionGroups.addAll(partitionGroups);
		return scheme.new Builder();
	}

	public final class Builder extends AbstractBuilder> {
		private Builder() {}

		public Builder withPartitionIdGetter(Function partitionIdGetter) {
			checkNotBuilt(this);
			RendezvousPartitionScheme.this.partitionIdGetter = partitionIdGetter;
			return this;
		}

		public Builder withCrdtProvider(Function> crdtProvider) {
			checkNotBuilt(this);
			RendezvousPartitionScheme.this.crdtProvider = crdtProvider;
			return this;
		}

		public Builder withRpcProvider(Function rpcProvider) {
			checkNotBuilt(this);
			RendezvousPartitionScheme.this.rpcProvider = rpcProvider;
			return this;
		}

		public Builder withPartitionGroup(RendezvousPartitionGroup
 partitionGroup) {
			checkNotBuilt(this);
			RendezvousPartitionScheme.this.partitionGroups.add(partitionGroup);
			return this;
		}

		public > Builder withKeyHashFn(ToIntFunction keyHashFn) {
			checkNotBuilt(this);
			RendezvousPartitionScheme.this.keyHashFn = keyHashFn;
			return this;
		}

		@Override
		protected RendezvousPartitionScheme
 doBuild() {
			return RendezvousPartitionScheme.this;
		}
	}

	@Override
	public Set
 getPartitions() {
		return partitionGroups.stream().flatMap(g -> g.getPartitionIds().stream()).collect(toSet());
	}

	@Override
	public ICrdtStorage provideCrdtConnection(P partition) {
		return crdtProvider.apply(partition);
	}

	@Override
	public RpcStrategy provideRpcConnection(P partition) {
		return rpcProvider.apply(partition);
	}

	@Override
	public > @Nullable Sharder createSharder(List
 alive) {
		Set
 aliveSet = new HashSet<>(alive);
		List> sharders = new ArrayList<>();
		for (RendezvousPartitionGroup
 partitionGroup : partitionGroups) {
			int deadPartitions = difference(partitionGroup.getPartitionIds(), aliveSet).size();

			if (partitionGroup.isRepartition()) {
				int aliveSize = partitionGroup.getPartitionIds().size() - deadPartitions;
				if (aliveSize < partitionGroup.getReplicaCount()) return null;
			} else if (deadPartitions != 0) return null;

			//noinspection unchecked
			RendezvousHashSharder sharder = RendezvousHashSharder.create(
				((ToIntFunction) keyHashFn),
				p -> partitionIdGetter.apply(p).hashCode(),
				partitionGroup.getPartitionIds(),
				alive,
				partitionGroup.getReplicaCount(), partitionGroup.isRepartition());
			sharders.add(sharder);
		}
		return RendezvousHashSharder.unionOf(sharders);
	}

	@Override
	public > RpcStrategy createRpcStrategy(Function keyGetter) {
		List rendezvousHashings = new ArrayList<>();
		for (RendezvousPartitionGroup
 partitionGroup : partitionGroups) {
			if (!partitionGroup.isActive()) continue;
			//noinspection unchecked
			rendezvousHashings.add(
				RendezvousHashing.builder(req ->
						((ToIntFunction) keyHashFn).applyAsInt(keyGetter.apply(req)))
					.withBuckets(NUMBER_OF_BUCKETS)
					.withHashBucketFunction((p, bucket) -> RendezvousHashSharder.hashBucket(partitionIdGetter.apply((P) p).hashCode(), bucket))
					.initialize(rendezvousHashing -> {
						for (P partitionId : partitionGroup.getPartitionIds()) {
							rendezvousHashing.withShard(partitionId, provideRpcConnection(partitionId));
						}
						if (!partitionGroup.isRepartition()) {
							rendezvousHashing.withReshardings(partitionGroup.getReplicaCount());
						}
					})
					.build());
		}
		final int count = rendezvousHashings.size();
		return RpcStrategies.sharding(item -> keyGetter.apply(item).hashCode() % count, rendezvousHashings);
	}

	@Override
	public boolean isReadValid(Collection
 alive) {
		Set
 aliveSet = new HashSet<>(alive);
		for (RendezvousPartitionGroup partitionGroup : partitionGroups) {
			int deadPartitions = difference(partitionGroup.getPartitionIds(), aliveSet).size();
			if (deadPartitions < partitionGroup.getReplicaCount()) {
				return true;
			}
		}
		return false;
	}

	@VisibleForTesting
	List> getPartitionGroups() {
		return partitionGroups;
	}
}