com.yandex.ydb.core.grpc.UnaryStreamToConsumer Maven / Gradle / Ivy

Go to download
package com.yandex.ydb.core.grpc;

import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.function.Consumer;

import com.yandex.ydb.core.Issue;
import com.yandex.ydb.core.Result;
import com.yandex.ydb.core.StatusCode;
import io.grpc.ClientCall;
import io.grpc.Metadata;
import io.grpc.Status;


/**
 * @author Sergey Polovko
 */
public class UnaryStreamToConsumer extends ClientCall.Listener {

    private static final AtomicIntegerFieldUpdater acceptedUpdater =
        AtomicIntegerFieldUpdater.newUpdater(UnaryStreamToConsumer.class, "accepted");

    private final Consumer> consumer;
    private volatile int accepted = 0;
    private T value;

    public UnaryStreamToConsumer(Consumer> consumer) {
        this.consumer = consumer;
    }

    @Override
    public void onMessage(T value) {
        if (this.value != null) {
            Issue issue = Issue.of("More than one value received for gRPC unary call", Issue.Severity.ERROR);
            accept(Result.fail(StatusCode.CLIENT_INTERNAL_ERROR, issue));
        } else {
            this.value = value;
        }
    }

    @Override
    public void onClose(Status status, Metadata trailers) {
        if (status.isOk()) {
            if (value == null) {
                Issue issue = Issue.of("No value received for gRPC unary call", Issue.Severity.ERROR);
                accept(Result.fail(StatusCode.CLIENT_INTERNAL_ERROR, issue));
            } else {
                accept(Result.success(value));
            }
        } else {
            accept(GrpcStatuses.toResult(status));
        }
    }

    void accept(Result result) {
        // protect from double accept
        if (acceptedUpdater.compareAndSet(this, 0, 1)) {
            consumer.accept(result);
        }
    }
}