com.linkedin.restli.client.DisruptRestClient Maven / Gradle / Ivy

Go to download
/*
   Copyright (c) 2017 LinkedIn Corp.

   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 com.linkedin.restli.client;

import com.linkedin.common.callback.Callback;
import com.linkedin.common.callback.Callbacks;
import com.linkedin.common.util.None;
import com.linkedin.r2.disruptor.DisruptContext;
import com.linkedin.r2.message.RequestContext;
import com.linkedin.restli.client.multiplexer.MultiplexedRequest;
import com.linkedin.restli.client.multiplexer.MultiplexedResponse;
import com.linkedin.restli.common.ResourceMethod;
import com.linkedin.restli.disruptor.DisruptRestController;
import com.linkedin.util.ArgumentUtil;


/**
 * Decorator Rest.li {@link Client} implementation that evaluates each {@link Request}
 * against a provided {@link DisruptRestController} instance and writes the evaluated
 * {@link DisruptContext} into the {@link RequestContext} object associated the request.
 * The #sendRequest operation is eventually delegated to the decorated Rest.li {@link Client}.
 *
 * @author Sean Sheng
 */
public class DisruptRestClient implements Client
{
  private final Client _client;
  private final DisruptRestController _controller;

  public DisruptRestClient(Client client, DisruptRestController controller)
  {
    ArgumentUtil.notNull(client, "client");
    ArgumentUtil.notNull(controller, "controller");

    _client = client;
    _controller = controller;
  }

  @Override
  public void shutdown(Callback callback)
  {
    _client.shutdown(callback);
  }

  @Override
  public  ResponseFuture sendRequest(Request request, RequestContext requestContext)
  {
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext);
  }

  @Override
  public  ResponseFuture sendRequest(Request request, RequestContext requestContext,
      ErrorHandlingBehavior errorHandlingBehavior)
  {
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext, errorHandlingBehavior);
  }

  @Override
  public  ResponseFuture sendRequest(RequestBuilder> requestBuilder,
      RequestContext requestContext)
  {
    Request request = requestBuilder.build();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext);
  }

  @Override
  public  ResponseFuture sendRequest(RequestBuilder> requestBuilder,
      RequestContext requestContext, ErrorHandlingBehavior errorHandlingBehavior)
  {
    Request request = requestBuilder.build();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext, errorHandlingBehavior);
  }

  @Override
  public  void sendRequest(Request request, RequestContext requestContext, Callback> callback)
  {
    doEvaluateDisruptContext(request, requestContext);
    _client.sendRequest(request, requestContext, callback);
  }

  @Override
  public  void sendRequest(RequestBuilder> requestBuilder, RequestContext requestContext,
      Callback> callback)
  {
    Request request = requestBuilder.build();
    doEvaluateDisruptContext(request, requestContext);
    _client.sendRequest(request, requestContext, callback);
  }

  @Override
  public  ResponseFuture sendRequest(Request request)
  {
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext);
  }

  @Override
  public  ResponseFuture sendRequest(Request request, ErrorHandlingBehavior errorHandlingBehavior)
  {
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext, errorHandlingBehavior);
  }

  @Override
  public  ResponseFuture sendRequest(RequestBuilder> requestBuilder)
  {
    Request request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext);
  }

  @Override
  public  ResponseFuture sendRequest(RequestBuilder> requestBuilder,
      ErrorHandlingBehavior errorHandlingBehavior)
  {
    Request request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    return _client.sendRequest(request, requestContext, errorHandlingBehavior);
  }

  @Override
  public  void sendRequest(Request request, Callback> callback)
  {
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    _client.sendRequest(request, requestContext, callback);
  }

  @Override
  public  void sendRequest(RequestBuilder> requestBuilder, Callback> callback)
  {
    Request request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(request, requestContext);
    _client.sendRequest(request, requestContext, callback);
  }

  @Override
  public void sendRequest(MultiplexedRequest multiplexedRequest)
  {
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(requestContext);
    _client.sendRequest(multiplexedRequest, requestContext, Callbacks.empty());
  }

  @Override
  public void sendRequest(MultiplexedRequest multiplexedRequest, Callback callback)
  {
    RequestContext requestContext = new RequestContext();
    doEvaluateDisruptContext(requestContext);
    _client.sendRequest(multiplexedRequest, requestContext, callback);
  }

  @Override
  public void sendRequest(MultiplexedRequest multiplexedRequest, RequestContext requestContext,
      Callback callback)
  {
    doEvaluateDisruptContext(requestContext);
    _client.sendRequest(multiplexedRequest, requestContext, callback);
  }

  /**
   * Evaluates if a {@link MultiplexedRequest} should be disrupted, against the {@link DisruptRestController}
   * and store the corresponding {@link DisruptContext} into the {@link RequestContext}. However,
   * if disrupt source is already set in the ReuqestContext, the method does not evaluate further.
   *
   * @param requestContext Context of the request
   * @param  Request template
   */
  private  void doEvaluateDisruptContext(RequestContext requestContext)
  {
    // Do not evaluate further if disrupt source key is already present in the request context
    if (isDisruptSourceAlreadySet(requestContext))
    {
      return;
    }

    requestContext.putLocalAttr(DisruptContext.DISRUPT_SOURCE_KEY, _controller.getClass().getCanonicalName());

    final DisruptContext disruptContext = _controller.getDisruptContext(MULTIPLEXER_RESOURCE);
    if (disruptContext != null)
    {
      requestContext.putLocalAttr(DisruptContext.DISRUPT_CONTEXT_KEY, disruptContext);
    }
  }

  /**
   * Evaluates if a {@link Request} should be disrupted, against the {@link DisruptRestController}
   * and store the corresponding {@link DisruptContext} into the {@link RequestContext}. However,
   * if disrupt source is already set in the ReuqestContext, the method does not evaluate further.
   *
   * @param request Request
   * @param requestContext Context associated with the request
   * @param  Request template
   */
  private  void doEvaluateDisruptContext(Request request, RequestContext requestContext)
  {
    // Do not evaluate further if disrupt source key is already present in the request context
    if (isDisruptSourceAlreadySet(requestContext))
    {
      return;
    }

    // Marks the disrupt source key property in the request context to indicate that disrupt has already been evaluated
    requestContext.putLocalAttr(DisruptContext.DISRUPT_SOURCE_KEY, _controller.getClass().getCanonicalName());

    final ResourceMethod method = request.getMethod();
    final String resource = request.getBaseUriTemplate();
    final String name = request.getMethodName();
    final DisruptContext disruptContext;
    if (name == null)
    {
      disruptContext = _controller.getDisruptContext(resource, method);
    }
    else
    {
      disruptContext = _controller.getDisruptContext(resource, method, name);
    }

    if (disruptContext != null)
    {
      requestContext.putLocalAttr(DisruptContext.DISRUPT_CONTEXT_KEY, disruptContext);
    }
  }

  /**
   * Evaluates if the disrupt source is already set in the {@link RequestContext}
   *
   * @param requestContext Context associated with the request
   * @return {@code true} if disrupt source is already set, {@code false} otherse;
   */
  private boolean isDisruptSourceAlreadySet(RequestContext requestContext)
  {
    return requestContext.getLocalAttr(DisruptContext.DISRUPT_SOURCE_KEY) != null;
  }
}