io.hekate.messaging.internal.MessageOperation Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hekate-core Show documentation
Java library for cluster communications and computing.
There is a newer version: 4.1.3
/*
 * Copyright 2020 The Hekate Project
 *
 * The Hekate Project licenses this file to you 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.hekate.messaging.internal;

import io.hekate.cluster.ClusterNodeId;
import io.hekate.messaging.intercept.OutboundType;
import io.hekate.messaging.loadbalance.LoadBalancerException;
import io.hekate.messaging.operation.ResponsePart;
import io.hekate.messaging.retry.FailedAttempt;
import io.hekate.messaging.retry.RetryBackoffPolicy;
import io.hekate.messaging.retry.RetryCallback;
import io.hekate.messaging.retry.RetryCondition;
import io.hekate.messaging.retry.RetryErrorPredicate;
import io.hekate.messaging.retry.RetryRoutingPolicy;
import io.hekate.partition.PartitionMapper;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;

import static java.util.concurrent.atomic.AtomicIntegerFieldUpdater.newUpdater;

abstract class MessageOperation {
    private static final int STATE_PENDING = 0;

    private static final int STATE_COMPLETED = 1;

    private static final AtomicIntegerFieldUpdater STATE = newUpdater(MessageOperation.class, "state");

    private final T message;

    private final MessageOperationOpts opts;

    private final Object affinityKey;

    private final int affinity;

    private final RetryErrorPredicate retryErr;

    private final RetryCondition retryCondition;

    private final RetryBackoffPolicy retryBackoff;

    private final RetryCallback retryCallback;

    private final RetryRoutingPolicy retryRoute;

    private final int maxAttempts;

    private final long timeout;

    private final MessagingWorker worker;

    private final MessagingGatewayContext gateway;

    private Future timeoutFuture;

    private SendPressureGuard sendPressure;

    @SuppressWarnings("unused") // <-- Updated via AtomicIntegerFieldUpdater.
    private volatile int state;

    public MessageOperation(
        T message,
        Object affinityKey,
        long timeout,
        int maxAttempts,
        RetryErrorPredicate retryErr,
        RetryCondition retryCondition,
        RetryBackoffPolicy retryBackoff,
        RetryCallback retryCallback,
        RetryRoutingPolicy retryRoute,
        MessagingGatewayContext gateway,
        MessageOperationOpts opts,
        boolean threadAffinity
    ) {
        this.message = message;
        this.affinityKey = affinityKey;
        this.maxAttempts = maxAttempts;
        this.timeout = timeout;
        this.gateway = gateway;
        this.retryErr = retryErr;
        this.retryCondition = retryCondition;
        this.retryBackoff = retryBackoff;
        this.retryCallback = retryCallback;
        this.retryRoute = retryRoute;
        this.opts = opts;

        if (affinityKey == null) {
            // Use artificial affinity.
            affinity = ThreadLocalRandom.current().nextInt();

            if (threadAffinity) {
                worker = gateway.async().workerFor(affinity);
            } else {
                worker = gateway.async().pooledWorker();
            }
        } else {
            // Use key-based affinity.
            affinity = affinityKey.hashCode();

            worker = gateway.async().workerFor(affinity);
        }
    }

    public abstract ClusterNodeId route(PartitionMapper mapper, Optional prevFailure) throws LoadBalancerException;

    public abstract OutboundType type();

    public abstract boolean shouldRetry(ResponsePart response);

    public abstract CompletableFuture future();

    protected abstract void doReceiveFinal(ResponsePart response);

    protected abstract void doFail(Throwable error);

    public long timeout() {
        return timeout;
    }

    public boolean hasTimeout() {
        return timeout > 0;
    }

    public void registerTimeout(Future timeoutFuture) {
        this.timeoutFuture = timeoutFuture;
    }

    public void registerSendPressure(SendPressureGuard sendPressure) {
        this.sendPressure = sendPressure;
    }

    public boolean isDone() {
        return state == STATE_COMPLETED;
    }

    public boolean canRetry() {
        return retryCondition == null || retryCondition.shouldRetry();
    }

    public void onRetry(FailedAttempt failure) {
        if (retryCallback != null) {
            retryCallback.onRetry(failure);
        }
    }

    public boolean complete(Throwable error, ResponsePart response) {
        if (response != null && !response.isLastPart()) {
            // Do not complete on partial responses.
            doReceivePartial(response);
        } else if (STATE.compareAndSet(this, STATE_PENDING, STATE_COMPLETED)) {
            if (sendPressure != null) {
                sendPressure.onDequeue();
            }

            if (timeoutFuture != null) {
                timeoutFuture.cancel(false);
            }

            if (error == null) {
                doReceiveFinal(response);
            } else {
                doFail(error);
            }

            return true;
        }

        return false;
    }

    public boolean shouldExpireOnTimeout() {
        return true;
    }

    public T message() {
        return message;
    }

    public RetryErrorPredicate retryErrorPolicy() {
        return retryErr;
    }

    public RetryRoutingPolicy retryRoute() {
        return retryRoute;
    }

    public int maxAttempts() {
        return maxAttempts;
    }

    public MessagingGatewayContext gateway() {
        return gateway;
    }

    public MessageOperationOpts opts() {
        return opts;
    }

    public Object affinityKey() {
        return affinityKey;
    }

    public boolean hasAffinity() {
        return affinityKey != null;
    }

    public int affinity() {
        return affinity;
    }

    public MessagingWorker worker() {
        return worker;
    }

    public RetryBackoffPolicy retryBackoff() {
        return retryBackoff;
    }

    protected void doReceivePartial(ResponsePart response) {
        throw new UnsupportedOperationException(getClass().getSimpleName() + " can't receive " + response);
    }
}