com.netflix.concurrency.limits.grpc.client.ConcurrencyLimitClientInterceptor Maven / Gradle / Ivy

Go to download
package com.netflix.concurrency.limits.grpc.client;

import com.google.common.base.Preconditions;
import com.netflix.concurrency.limits.Limiter;

import io.grpc.CallOptions;
import io.grpc.Channel;
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.ForwardingClientCall;
import io.grpc.ForwardingClientCallListener;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.Status;
import io.grpc.Status.Code;

import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.Nullable;

/**
 * ClientInterceptor that enforces per service and/or per method concurrent request limits and returns
 * a Status.UNAVAILABLE when that limit has been reached.  
 */
public class ConcurrencyLimitClientInterceptor implements ClientInterceptor {
    private static final Status LIMIT_EXCEEDED_STATUS = Status.UNAVAILABLE.withDescription("Concurrency limit reached");
    
    private final Limiter grpcLimiter;
    
    public ConcurrencyLimitClientInterceptor(final Limiter grpcLimiter) {
        Preconditions.checkArgument(grpcLimiter != null, "GrpcLimiter cannot not be null");
        this.grpcLimiter = grpcLimiter;
    }
    
    @Override
    public  ClientCall interceptCall(final MethodDescriptor method,
            final CallOptions callOptions, final Channel next) {
        final Optional listener = grpcLimiter.acquire(new GrpcClientRequestContext() {
            @Override
            public MethodDescriptor getMethod() {
                return method;
            }

            @Override
            public CallOptions getCallOptions() {
                return callOptions;
            }
        });
        
        if (!listener.isPresent()) {
            return new ClientCall() {
                @Override
                public void start(io.grpc.ClientCall.Listener responseListener, Metadata headers) {
                    responseListener.onClose(LIMIT_EXCEEDED_STATUS, new Metadata());
                }

                @Override
                public void request(int numMessages) {
                }

                @Override
                public void cancel(String message, Throwable cause) {
                }

                @Override
                public void halfClose() {
                }

                @Override
                public void sendMessage(ReqT message) {
                }
            };
        }
        
        // Perform the operation and release the limiter once done.  
        return new ForwardingClientCall.SimpleForwardingClientCall(next.newCall(method, callOptions)) {
            final AtomicBoolean done = new AtomicBoolean(false);
            
            @Override
            public void start(final Listener responseListener, final Metadata headers) {
                super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener(responseListener) {
                    @Override
                    public void onClose(final Status status, final Metadata trailers) {
                        try {
                            super.onClose(status, trailers);
                        } finally {
                            listener.ifPresent(l -> {
                                if (done.compareAndSet(false, true)) {
                                    if (status.isOk()) {
                                        l.onSuccess();
                                    } else if (Code.UNAVAILABLE == status.getCode()) {
                                        l.onDropped();
                                    } else {
                                        l.onIgnore();
                                    }
                                }
                            });
                        }
                    }
                }, headers);
            }
            
            @Override
            public void cancel(final @Nullable String message, final @Nullable Throwable cause) {
                try {
                    super.cancel(message, cause);
                } finally {
                    if (done.compareAndSet(false, true)) {
                        listener.ifPresent(Limiter.Listener::onIgnore);
                    }
                }
            }
        };
    }
}