io.activej.etcd.EtcdUtils Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of activej-etcd Show documentation
Etcd Codecs, and transactional support for ActiveJ project.
The newest version!
package io.activej.etcd;

import io.activej.common.time.CurrentTimeProvider;
import io.activej.etcd.codec.key.EtcdKeyEncoder;
import io.activej.etcd.codec.kv.EtcdKVDecoder;
import io.activej.etcd.codec.kv.EtcdKVEncoder;
import io.activej.etcd.codec.kv.KeyValue;
import io.activej.etcd.codec.value.EtcdValueCodec;
import io.activej.etcd.codec.value.EtcdValueCodecs;
import io.activej.etcd.codec.value.EtcdValueEncoder;
import io.activej.etcd.exception.EtcdException;
import io.activej.etcd.exception.MalformedEtcdDataException;
import io.activej.etcd.exception.NoKeyFoundException;
import io.activej.etcd.exception.TransactionNotSucceededException;
import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.KV;
import io.etcd.jetcd.Response;
import io.etcd.jetcd.Watch;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.kv.TxnResponse;
import io.etcd.jetcd.op.Cmp;
import io.etcd.jetcd.op.CmpTarget;
import io.etcd.jetcd.op.Op;
import io.etcd.jetcd.options.DeleteOption;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.options.WatchOption;
import io.etcd.jetcd.watch.WatchEvent;
import io.etcd.jetcd.watch.WatchResponse;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;

import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.CompletionStage;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.UnaryOperator;
import java.util.stream.Collector;

import static io.activej.common.Checks.checkArgument;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.concurrent.CompletableFuture.completedFuture;
import static java.util.concurrent.CompletableFuture.failedFuture;

public class EtcdUtils {

	public static final EtcdValueCodec TOUCH_TIMESTAMP_CODEC = EtcdValueCodecs.ofLongString();

	private static final String DEADLINE_EXCEEDED_STATUS_MESSAGE = "context deadline exceeded";

	public static ByteSequence byteSequenceFrom(char ch) {
		checkArgument(ch <= 0x7f);
		return byteSequenceFrom((byte) ch);
	}

	public static ByteSequence byteSequenceFrom(byte b) {
		return ByteSequence.from(new byte[]{b});
	}

	public static ByteSequence byteSequenceFrom(String string) {
		return ByteSequence.from(string, UTF_8);
	}

	public record CheckoutResponse(io.etcd.jetcd.Response.Header header, T response) {}

	public record CheckoutRequest(ByteSequence prefix, EtcdKVDecoder codec, Collector collector) {

		public static  CheckoutRequest, R> ofMapEntry(
			ByteSequence prefix, EtcdKVDecoder> codec, Collector, ?, ? extends R> collector
		) {
			//noinspection unchecked,rawtypes
			return CheckoutRequest., R>of(prefix, codec, (Collector) collector);
		}

		public static  CheckoutRequest of(
			ByteSequence prefix, EtcdKVDecoder codec, Collector collector
		) {
			return new CheckoutRequest<>(prefix, codec, collector);
		}

	}

	public static  CompletableFuture> checkout(KV client, long revision,
		ByteSequence prefix, EtcdKVDecoder codec, Collector collector
	) {
		//noinspection unchecked
		return checkout(client, revision, new CheckoutRequest[]{new CheckoutRequest<>(prefix, codec, collector)},
			(header, objects) -> new CheckoutResponse<>(header, (R) objects[0]));
	}

	public interface CheckoutFinisher {
		T finish(Response.Header header, Object[] objects) throws MalformedEtcdDataException;
	}

	public static  CompletableFuture checkout(KV client, long revision,
		CheckoutRequest[] checkoutRequests, CheckoutFinisher checkoutFinisher
	) {
		return client.txn()
			.Then(Arrays.stream(checkoutRequests)
				.map(checkoutRequest -> Op.get(
					checkoutRequest.prefix,
					GetOption.builder()
						.isPrefix(true)
						.withRevision(revision)
						.build()))
				.toArray(Op[]::new))
			.commit()
			.exceptionallyCompose(e -> failedFuture(new EtcdException("Checkout failed", convertStatusException(e.getCause()))))
			.thenCompose(new Function>() {
				@Override
				public CompletionStage apply(TxnResponse response) {
					try {
						Object[] result = new Object[checkoutRequests.length];
						for (int i = 0; i < checkoutRequests.length; i++) {
							//noinspection unchecked
							CheckoutRequest checkoutRequest = (CheckoutRequest) checkoutRequests[i];
							GetResponse getResponse = response.getGetResponses().get(i);
							result[i] = getCheckoutResult(checkoutRequest.prefix, checkoutRequest.codec, checkoutRequest.collector, getResponse);
						}
						return completedFuture(checkoutFinisher.finish(response.getHeader(), result));
					} catch (MalformedEtcdDataException e) {
						return failedFuture(e);
					}
				}

				static  R getCheckoutResult(ByteSequence prefix, EtcdKVDecoder codec, Collector collector, GetResponse response) throws MalformedEtcdDataException {
					A accumulator = collector.supplier().get();
					BiConsumer accumulatorFn = collector.accumulator();
					for (var kv : response.getKvs()) {
						KV item;
						ByteSequence key = kv.getKey();
						try {
							item = codec.decodeKV(new KeyValue(key.substring(prefix.size()), kv.getValue()));
						} catch (MalformedEtcdDataException e) {
							throw new MalformedEtcdDataException("Failed to decode KV of key '" + key + '\'', e);
						}
						accumulatorFn.accept(accumulator, item);
					}
					return collector.finisher().apply(accumulator);
				}
			});
	}

	public static  Watch.Watcher watch(Watch watch, long revision,
		ByteSequence prefix, EtcdKVDecoder codec, EtcdEventProcessor eventProcessor,
		EtcdListener listener
	) {
		return watch(watch, revision,
			new WatchRequest[]{new WatchRequest<>(prefix, codec, eventProcessor)},
			new EtcdListener<>() {
				@Override
				public void onConnectionEstablished() {
					listener.onConnectionEstablished();
				}

				@Override
				public void onNext(long revision, Object[] operations) throws MalformedEtcdDataException {
					assert operations.length == 1;
					//noinspection unchecked
					listener.onNext(revision, (R) operations[0]);
				}

				@Override
				public void onError(Throwable throwable) {
					listener.onError(throwable);
				}

				@Override
				public void onCompleted() {
					listener.onCompleted();
				}
			});
	}

	public record WatchRequest(ByteSequence prefix, EtcdKVDecoder codec, EtcdEventProcessor eventProcessor) {

		public static  WatchRequest, R> ofMapEntry(
			ByteSequence prefix, EtcdKVDecoder> codec, EtcdEventProcessor, ? extends R> eventProcessor
		) {
			return of(prefix, codec, eventProcessor);
		}

		public static  WatchRequest of(
			ByteSequence prefix, EtcdKVDecoder codec, EtcdEventProcessor eventProcessor
		) {
			return new WatchRequest<>(prefix, codec, eventProcessor);
		}

	}

	public static Watch.Watcher watch(Watch client, long revision,
		WatchRequest[] watchRequests,
		EtcdListener listener
	) {
		byte[][] prefixes = Arrays.stream(watchRequests).map(WatchRequest::prefix).map(ByteSequence::getBytes).toArray(byte[][]::new);
		int rootSize = 0;
		LOOP:
		for (; rootSize < prefixes[0].length; rootSize++) {
			for (int i = 1; i < prefixes.length; i++) {
				if (rootSize > prefixes[i].length || prefixes[0][rootSize] != prefixes[i][rootSize]) {
					break LOOP;
				}
			}
		}
		ByteSequence root = ByteSequence.from(Arrays.copyOf(prefixes[0], rootSize));

		//noinspection rawtypes,unchecked
		WatchRequest[] requests = Arrays.stream(watchRequests)
			.map(w -> new WatchRequest(w.prefix.substring(root.size()), w.codec, w.eventProcessor))
			.toArray(WatchRequest[]::new);

		return client.watch(
			root,
			WatchOption.builder()
				.isPrefix(true)
				.withRevision(revision)
				.withCreateNotify(true)
				.build(),
			new Watch.Listener() {
				long currentRevision = -1;

				@Override
				public void onNext(WatchResponse response) {
					if (response.isCreatedNotify()) {
						listener.onConnectionEstablished();
					}
					try {
						Object[] accumulators = new Object[requests.length];

						boolean firstRun = true;
						for (var event : response.getEvents()) {
							var keyValue = event.getKeyValue();
							var rootKey = keyValue.getKey().substring(root.size());

							long modRevision = keyValue.getModRevision();
							if (modRevision != currentRevision) {
								if (!firstRun) listener.onNext(currentRevision, accumulators);

								for (int i = 0; i < accumulators.length; i++) {
									accumulators[i] = requests[i].eventProcessor.createEventsAccumulator();
								}
								currentRevision = modRevision;
							}

							for (int i = 0; i < accumulators.length; i++) {
								var request = requests[i];
								if (!rootKey.startsWith(request.prefix)) continue;
								var key = rootKey.substring(request.prefix.size());
								if (event.getEventType() == WatchEvent.EventType.PUT) {
									Object kv;
									try {
										kv = request.codec.decodeKV(new KeyValue(key, keyValue.getValue()));
									} catch (MalformedEtcdDataException e) {
										throw new MalformedEtcdDataException("Failed to decode KV of key '" + rootKey + '\'', e);
									}
									//noinspection unchecked
									request.eventProcessor.onPut(accumulators[i], kv);
								}
								if (event.getEventType() == WatchEvent.EventType.DELETE) {
									Object k;
									try {
										k = request.codec.decodeKey(key);
									} catch (MalformedEtcdDataException e) {
										throw new MalformedEtcdDataException("Failed to decode key '" + rootKey + '\'', e);
									}
									//noinspection unchecked
									request.eventProcessor.onDelete(accumulators[i], k);
								}
							}
							firstRun = false;
						}
						if (currentRevision != -1) {
							listener.onNext(currentRevision, accumulators);
						}
					} catch (MalformedEtcdDataException e) {
						onError(e);
					}
				}

				@Override
				public void onError(Throwable throwable) {
					listener.onError(throwable);
				}

				@Override
				public void onCompleted() {
					listener.onCompleted();
				}
			});
	}

	public static void touch(TxnOps txn, ByteSequence key) {
		txn.put(key, ByteSequence.EMPTY, PutOption.DEFAULT);
	}

	public static void touchTimestamp(TxnOps txn, ByteSequence key, CurrentTimeProvider timeProvider) {
		ByteSequence value = TOUCH_TIMESTAMP_CODEC.encodeValue(timeProvider.currentTimeMillis());
		txn.put(key, value, PutOption.DEFAULT);
	}

	public record AtomicUpdateResponse(Response.Header header, T prevValue, T newValue) {}

	public static CompletableFuture> atomicAdd(KV client, ByteSequence key, int delta) {
		return atomicUpdate(client, key, EtcdValueCodecs.ofIntegerString(), value -> value + delta);
	}

	public static CompletableFuture> atomicAdd(KV client, ByteSequence key, long delta) {
		return atomicUpdate(client, key, EtcdValueCodecs.ofLongString(), value -> value + delta);
	}

	public static  CompletableFuture> atomicUpdate(
		KV client, ByteSequence key, EtcdValueCodec codec, UnaryOperator operator
	) {
		CompletableFuture> future = new CompletableFuture<>();
		atomicUpdate1(client, key, codec, operator, future);
		return future;
	}

	private static  void atomicUpdate1(
		KV client, ByteSequence key, EtcdValueCodec codec, UnaryOperator operator,
		CompletableFuture> future
	) {
		client
			.get(key, GetOption.builder().withSerializable(true).build())
			.whenComplete((getResponse, throwable) -> {
				if (throwable != null) {
					future.completeExceptionally(new EtcdException("Atomic update failed", convertStatusException(throwable.getCause())));
					return;
				}
				if (getResponse.getKvs().isEmpty()) {
					future.completeExceptionally(new NoKeyFoundException(key));
					return;
				}
				ByteSequence prevSequence = getResponse.getKvs().get(0).getValue();
				T prevValue;
				try {
					prevValue = codec.decodeValue(prevSequence);
				} catch (MalformedEtcdDataException e) {
					future.completeExceptionally(new MalformedEtcdDataException(
						"Failed to decode value of key '" + key + '\'', e));
					return;
				}
				T newValue = operator.apply(prevValue);
				ByteSequence newSequence = codec.encodeValue(newValue);
				atomicUpdate2(client, key, codec, operator, prevSequence, newSequence, prevValue, newValue, future);
			});
	}

	private static  void atomicUpdate2(
		KV client, ByteSequence key, EtcdValueCodec codec, UnaryOperator operator,
		ByteSequence prevSequence, ByteSequence newSequence, T prevValue, T newValue,
		CompletableFuture> future
	) {
		client.txn()
			.If(new Cmp(key, Cmp.Op.EQUAL, CmpTarget.value(prevSequence)))
			.Then(Op.put(key, newSequence, PutOption.DEFAULT))
			.commit()
			.whenComplete((txnResponse, throwable) -> {
				if (throwable != null) {
					future.completeExceptionally(new EtcdException("Atomic update failed", convertStatusException(throwable.getCause())));
					return;
				}
				if (!txnResponse.isSucceeded()) {
					atomicUpdate1(client, key, codec, operator, future);
					return;
				}
				future.complete(new AtomicUpdateResponse<>(txnResponse.getHeader(), prevValue, newValue));
			});
	}

	public static  void checkAndUpdate(TxnOps txn, ByteSequence key, EtcdValueEncoder codec, V prev, V next) {
		checkAndReplace(txn, key, codec.encodeValue(prev), codec.encodeValue(next));
	}

	public static  void checkAndUpdate(TxnOps txn, EtcdKVEncoder> codec, Map prev, Map next) {
		for (var prevEntry : prev.entrySet()) {
			K k = prevEntry.getKey();
			KeyValue prevKV = codec.encodeKV(prevEntry);
			if (next.containsKey(k)) {
				V nextV = next.get(k);
				KeyValue nextKV = codec.encodeKV(Map.entry(k, nextV));
				if (prevKV.key().equals(nextKV.key())) {
					checkAndReplace(txn, prevKV.key(), prevKV.value(), nextKV.value());
				} else {
					checkAndMove(txn, prevKV.key(), prevKV.value(), nextKV.key(), nextKV.value());
				}
			} else {
				checkAndDelete(txn, prevKV.key(), prevKV.value());
			}
		}
		for (var nextEntry : next.entrySet()) {
			if (prev.containsKey(nextEntry.getKey())) return;
			KeyValue nextKV = codec.encodeKV(nextEntry);
			checkAndInsert(txn, nextKV.key(), nextKV.value());
		}
	}

	public static  void checkAndInsert(TxnOps txn, EtcdKVEncoder> codec, Map next) {
		for (var nextEntry : next.entrySet()) {
			K k = nextEntry.getKey();
			V nextV = nextEntry.getValue();
			KeyValue nextKeyValue = codec.encodeKV(Map.entry(k, nextV));
			checkAndInsert(txn, nextKeyValue.key(), nextKeyValue.value());
		}
	}

	public static  void checkAndInsert(TxnOps txn, EtcdKVEncoder codec, Collection next) {
		for (var item : next) {
			KeyValue nextKV = codec.encodeKV(item);
			checkAndInsert(txn, nextKV.key(), nextKV.value());
		}
	}

	public static  void checkAndReplace(TxnOps txn, EtcdKVEncoder codec, Collection next) {
		for (var item : next) {
			KeyValue nextKV = codec.encodeKV(item);
			checkAndReplace(txn, nextKV.key(), nextKV.value());
		}
	}

	public static  void checkAndDelete(TxnOps txn, EtcdKVEncoder codec, Collection prev) {
		for (var item : prev) {
			KeyValue prevKV = codec.encodeKV(item);
			checkAndDelete(txn, prevKV.key(), prevKV.value());
		}
	}

	public static  void checkAndDelete(TxnOps txn, EtcdKeyEncoder codec, Collection prev) {
		for (var item : prev) {
			checkAndDelete(txn, codec.encodeKey(item));
		}
	}

	public static void checkAndReplace(TxnOps txn, ByteSequence key, ByteSequence prevValue, ByteSequence nextValue) {
		txn.cmp(key, Cmp.Op.EQUAL, CmpTarget.value(prevValue));
		txn.put(key, nextValue, PutOption.DEFAULT);
	}

	public static void checkAndReplace(TxnOps txn, ByteSequence key, ByteSequence nextValue) {
		txn.cmp(key, Cmp.Op.GREATER, CmpTarget.createRevision(0));
		txn.put(key, nextValue, PutOption.DEFAULT);
	}

	public static void checkAndMove(TxnOps txn, ByteSequence prevKey, ByteSequence prevValue, ByteSequence nextKey, ByteSequence nextValue) {
		txn.cmp(prevKey, Cmp.Op.EQUAL, CmpTarget.value(prevValue));
		checkAndMove(txn, prevKey, nextKey, nextValue);
	}

	public static void checkAndMove(TxnOps txn, ByteSequence prevKey, ByteSequence nextKey, ByteSequence nextValue) {
		txn.cmp(nextKey, Cmp.Op.EQUAL, CmpTarget.createRevision(0));
		txn.delete(prevKey, DeleteOption.DEFAULT);
		txn.put(nextKey, nextValue, PutOption.DEFAULT);
	}

	public static void checkAndInsert(TxnOps txn, ByteSequence key, ByteSequence value) {
		txn.cmp(key, Cmp.Op.EQUAL, CmpTarget.createRevision(0));
		txn.put(key, value, PutOption.DEFAULT);
	}

	public static void checkAndDelete(TxnOps txn, ByteSequence key, ByteSequence value) {
		txn.cmp(key, Cmp.Op.EQUAL, CmpTarget.value(value));
		checkAndDelete(txn, key);
	}

	public static void checkAndDelete(TxnOps txn, ByteSequence key) {
		txn.cmp(key, Cmp.Op.GREATER, CmpTarget.createRevision(0));
		txn.delete(key, DeleteOption.DEFAULT);
	}

	public static CompletableFuture executeTxnOps(KV client, ByteSequence prefix, Consumer txnOpsConsumer) {
		TxnOps txnOps = new TxnOps(prefix);
		txnOpsConsumer.accept(txnOps);
		return executeTxnOps(client, txnOps);
	}

	public static CompletableFuture executeTxnOps(KV client, TxnOps txnOps) {
		return client.txn()
			.If(txnOps.cmps.toArray(Cmp[]::new))
			.Then(txnOps.ops.toArray(Op[]::new))
			.commit()
			.exceptionallyCompose(e -> failedFuture(new EtcdException("Transaction failed", convertStatusException(e.getCause()))))
			.thenCompose(txnResponse ->
				txnResponse.isSucceeded() ?
					completedFuture(txnResponse) :
					failedFuture(new TransactionNotSucceededException()));
	}

	public static  CompletionStage convertStatusExceptionStage(Throwable throwable) {
		if (throwable instanceof CompletionException) {
			throwable = throwable.getCause();
		}
		return failedFuture(convertStatusException(throwable));
	}

	public static Throwable convertStatusException(Throwable throwable) {
		if (!(throwable instanceof StatusRuntimeException sre)) {
			return throwable;
		}
		Status status = Status.fromThrowable(throwable);
		return switch (status.getCode()) {
			case OK,
				CANCELLED,
				NOT_FOUND,
				ALREADY_EXISTS,
				DEADLINE_EXCEEDED,
				UNAVAILABLE,
				ABORTED -> status.asException(sre.getTrailers());
			case PERMISSION_DENIED,
				RESOURCE_EXHAUSTED,
				INVALID_ARGUMENT,
				OUT_OF_RANGE,
				FAILED_PRECONDITION,
				UNIMPLEMENTED,
				INTERNAL,
				DATA_LOSS,
				UNAUTHENTICATED -> status.asRuntimeException(sre.getTrailers());
			case UNKNOWN -> DEADLINE_EXCEEDED_STATUS_MESSAGE.equals(status.getDescription()) ?
				status.asException(sre.getTrailers()) :
				status.asRuntimeException(sre.getTrailers());
		};
	}

}