io.activej.rpc.client.sender.strategy.impl.RendezvousHashing Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of activej-rpc Show documentation
High-performance and fault-tolerant Remote Procedure Call module for building distributed applications. Provides a high-performance asynchronous binary RPC streaming protocol.
There is a newer version: 6.0-rc2
Show newest version
/*
 * Copyright (C) 2020 ActiveJ LLC.
 *
 * Licensed 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 io.activej.rpc.client.sender.strategy.impl;

import io.activej.async.callback.Callback;
import io.activej.common.HashUtils;
import io.activej.common.annotation.ExposedInternals;
import io.activej.common.builder.AbstractBuilder;
import io.activej.rpc.client.RpcClientConnectionPool;
import io.activej.rpc.client.sender.RpcSender;
import io.activej.rpc.client.sender.strategy.RpcStrategy;
import io.activej.rpc.protocol.RpcException;
import org.jetbrains.annotations.Nullable;

import java.net.InetSocketAddress;
import java.util.*;
import java.util.function.ToIntFunction;
import java.util.function.ToLongBiFunction;

import static io.activej.common.Checks.checkArgument;
import static io.activej.rpc.client.sender.strategy.RpcStrategies.server;
import static java.lang.Math.min;

@ExposedInternals
public final class RendezvousHashing implements RpcStrategy {
	public static final int DEFAULT_BUCKET_CAPACITY = 2048;
	public static final ToLongBiFunction DEFAULT_HASH_BUCKET_FN = (shardId, bucketN) ->
		(int) HashUtils.murmur3hash(((long) shardId.hashCode() << 32) | (bucketN & 0xFFFFFFFFL));
	public static final int DEFAULT_MIN_ACTIVE_SHARDS = 1;
	public static final int DEFAULT_MAX_RESHARDINGS = Integer.MAX_VALUE;

	public final ToIntFunction hashFn;
	public final Map shards;
	public ToLongBiFunction hashBucketFn;
	public int buckets;
	public int minActiveShards;
	public int reshardings;

	public RendezvousHashing(
		ToIntFunction hashFn,
		Map shards,
		ToLongBiFunction hashBucketFn,
		int buckets,
		int minActiveShards,
		int reshardings
	) {
		this.hashFn = hashFn;
		this.shards = shards;
		this.hashBucketFn = hashBucketFn;
		this.buckets = buckets;
		this.minActiveShards = minActiveShards;
		this.reshardings = reshardings;
	}

	public static  Builder builder(ToIntFunction hashFn) {
		return new RendezvousHashing(
			hashFn, new HashMap<>(), DEFAULT_HASH_BUCKET_FN,
			DEFAULT_BUCKET_CAPACITY, DEFAULT_MIN_ACTIVE_SHARDS, DEFAULT_MAX_RESHARDINGS).new Builder();
	}

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

		public Builder withHashBucketFunction(ToLongBiFunction hashBucketFn) {
			checkNotBuilt(this);
			RendezvousHashing.this.hashBucketFn = hashBucketFn;
			return this;
		}

		public Builder withBuckets(int buckets) {
			checkNotBuilt(this);
			checkArgument((buckets & (buckets - 1)) == 0);
			RendezvousHashing.this.buckets = buckets;
			return this;
		}

		public Builder withMinActiveShards(int minActiveShards) {
			checkNotBuilt(this);
			RendezvousHashing.this.minActiveShards = minActiveShards;
			return this;
		}

		public Builder withReshardings(int reshardings) {
			checkNotBuilt(this);
			RendezvousHashing.this.reshardings = reshardings;
			return this;
		}

		public Builder withShard(Object shardId, RpcStrategy strategy) {
			checkNotBuilt(this);
			shards.put(shardId, strategy);
			return this;
		}

		public Builder withShards(InetSocketAddress... addresses) {
			return withShards(List.of(addresses));
		}

		public Builder withShards(List addresses) {
			checkNotBuilt(this);
			for (InetSocketAddress address : addresses) {
				shards.put(address, server(address));
			}
			return this;
		}

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

	@Override
	public Set getAddresses() {
		HashSet result = new HashSet<>();
		for (RpcStrategy strategy : shards.values()) {
			result.addAll(strategy.getAddresses());
		}
		return result;
	}

	@Override
	public @Nullable RpcSender createSender(RpcClientConnectionPool pool) {
		Map shardsSenders = new HashMap<>();
		int activeShards = 0;
		for (Map.Entry entry : shards.entrySet()) {
			Object shardId = entry.getKey();
			RpcStrategy strategy = entry.getValue();
			RpcSender sender = strategy.createSender(pool);
			if (sender != null) {
				activeShards++;
			}
			shardsSenders.put(shardId, sender);
		}

		if (activeShards < minActiveShards) {
			return null;
		}

		RpcSender[] sendersBuckets = new RpcSender[buckets];

		//noinspection NullableProblems
		class ShardIdAndSender {
			final Object shardId;
			final @Nullable RpcSender rpcSender;
			private long hash;

			public long getHash() {
				return hash;
			}

			ShardIdAndSender(Object shardId, RpcSender rpcSender) {
				this.shardId = shardId;
				this.rpcSender = rpcSender;
			}
		}

		ShardIdAndSender[] toSort = new ShardIdAndSender[shards.size()];
		int i = 0;
		for (Map.Entry entry : shardsSenders.entrySet()) {
			toSort[i++] = new ShardIdAndSender(entry.getKey(), entry.getValue());
		}

		for (int bucket = 0; bucket < sendersBuckets.length; bucket++) {
			for (ShardIdAndSender obj : toSort) {
				obj.hash = hashBucketFn.applyAsLong(obj.shardId, bucket);
			}

			Arrays.sort(toSort, Comparator.comparingLong(ShardIdAndSender::getHash).reversed());

			for (int j = 0; j < min(shards.size(), reshardings); j++) {
				RpcSender rpcSender = toSort[j].rpcSender;
				if (rpcSender != null) {
					sendersBuckets[bucket] = rpcSender;
					break;
				}
			}
		}

		return new Sender(hashFn, sendersBuckets);
	}

	public static final class Sender implements RpcSender {
		private final ToIntFunction