All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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

Go to download

Pegasus is a framework for building robust, scalable service architectures using dynamic discovery and simple asychronous type-checked REST + JSON APIs.

There is a newer version: 5.0.19
Show newest version
/*
   Copyright (c) 2012 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.
*/

/**
 * $Id: $
 */

package com.linkedin.restli.client;


import com.linkedin.common.callback.Callback;
import com.linkedin.d2.balancer.KeyMapper;
import com.linkedin.d2.balancer.ServiceUnavailableException;
import com.linkedin.d2.balancer.util.MapKeyResult;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.r2.message.RequestContext;
import com.linkedin.restli.client.response.BatchKVResponse;
import com.linkedin.restli.common.BatchResponse;
import com.linkedin.restli.common.EntityResponse;
import com.linkedin.restli.common.RestConstants;
import com.linkedin.restli.common.TypeSpec;
import com.linkedin.restli.common.UpdateStatus;

import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * @author Josh Walker
 * @version $Revision: $
 */

public class ScatterGatherBuilder
{
  private final KeyMapper _mapper;
  private final String D2_URI_PREFIX = "d2://";

  public ScatterGatherBuilder(KeyMapper mapper)
  {
    _mapper = mapper;
  }

  // return value contains the request info and the unmapped keys (also the cause)
  // V2 is here to differentiate it from the older API
  @SuppressWarnings("deprecation")
  public ScatterGatherResult buildRequestsV2(BatchGetRequest request, RequestContext requestContext)
    throws ServiceUnavailableException
  {
    Set idObjects = request.getObjectIds();

    MapKeyResult mapKeyResult = mapKeys(request, idObjects);

    Map> batches = mapKeyResult.getMapResult();
    Collection> scatterGatherRequests = new ArrayList>(batches.size());

    for (Map.Entry> batch : batches.entrySet())
    {
      BatchGetRequestBuilder builder = new BatchGetRequestBuilder(request.getBaseUriTemplate(),
                                                                                        request.getResponseDecoder(),
                                                                                        request.getResourceSpec(),
                                                                                        request.getRequestOptions());
      builder.ids(batch.getValue());
      for (Map.Entry param : request.getQueryParamsObjects().entrySet())
      {
        if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM))
        {
          // keep all non-batch query parameters since they could be request specific
          builder.setParam(param.getKey(), param.getValue());
        }
      }
      for (Map.Entry header : request.getHeaders().entrySet())
      {
        builder.setHeader(header.getKey(), header.getValue());
      }

      RequestContext context = requestContext.clone();
      KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());

      scatterGatherRequests.add(new RequestInfo(builder.build(), context));
    }

    return new ScatterGatherResult(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
  }

  @SuppressWarnings("deprecation")
  public  KVScatterGatherResult> buildRequests(BatchGetEntityRequest request, RequestContext requestContext)
    throws ServiceUnavailableException
  {
    @SuppressWarnings("unchecked")
    final Set idObjects = (Set) request.getObjectIds();

    final MapKeyResult mapKeyResult = mapKeys(request, idObjects);

    final Map> batches = mapKeyResult.getMapResult();
    final Collection>> scatterGatherRequests = new ArrayList>>(batches.size());

    for (Map.Entry> batch : batches.entrySet())
    {
      final BatchGetEntityRequestBuilder builder = new BatchGetEntityRequestBuilder(request.getBaseUriTemplate(),
                                                                                                request.getResponseDecoder(),
                                                                                                request.getResourceSpec(),
                                                                                                request.getRequestOptions());
      builder.ids(batch.getValue());
      for (Map.Entry param : request.getQueryParamsObjects().entrySet())
      {
        if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM))
        {
          // keep all non-batch query parameters since they could be request specific
          builder.setParam(param.getKey(), param.getValue());
        }
      }
      for (Map.Entry header : request.getHeaders().entrySet())
      {
        builder.setHeader(header.getKey(), header.getValue());
      }

      final RequestContext context = requestContext.clone();
      KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());

      scatterGatherRequests.add(new KVRequestInfo>(builder.build(), context));
    }

    return new KVScatterGatherResult>(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
  }

  @SuppressWarnings({ "unchecked", "deprecation" })
  public  KVScatterGatherResult buildRequestsKV(BatchGetKVRequest request, RequestContext requestContext)
      throws ServiceUnavailableException
  {
    Set idObjects = (Set) request.getObjectIds();

    MapKeyResult mapKeyResult = mapKeys(request, idObjects);

    Map> batches = mapKeyResult.getMapResult();
    Collection> scatterGatherRequests = new ArrayList>(batches.size());

    for (Map.Entry> batch : batches.entrySet())
    {
      BatchGetRequestBuilder builder =
          new BatchGetRequestBuilder(request.getBaseUriTemplate(),
            (Class)request.getResourceProperties().getValueType().getType(),
            request.getResourceSpec(),
            request.getRequestOptions());

      builder.ids(batch.getValue());
      for (Map.Entry param : request.getQueryParamsObjects().entrySet())
      {
        if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM))
        {
          // keep all non-batch query parameters since they could be request specific
          builder.setParam(param.getKey(), param.getValue());
        }
      }
      for (Map.Entry header : request.getHeaders().entrySet())
      {
        builder.setHeader(header.getKey(), header.getValue());
      }

      RequestContext context = requestContext.clone();
      KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());

      scatterGatherRequests.add(new KVRequestInfo(builder.buildKV(), context));
    }

    return new KVScatterGatherResult(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
  }

  @SuppressWarnings("deprecation")
  public  KVScatterGatherResult buildRequests(BatchUpdateRequest request, RequestContext requestContext)
    throws ServiceUnavailableException
  {
    Set idObjects = request.getObjectIds();
    Collection ids = new HashSet(idObjects.size());
    for (Object o : idObjects)
    {
      @SuppressWarnings("unchecked")
      K k = (K) o;
      ids.add(k);
    }

    MapKeyResult mapKeyResult = mapKeys(request, ids);

    @SuppressWarnings("unchecked")
    TypeSpec valueType = (TypeSpec) request.getResourceProperties().getValueType();
    Map> batches = keyMapToInput(mapKeyResult, request);
    Collection> scatterGatherRequests = new ArrayList>(batches.size());

    for (Map.Entry> batch : batches.entrySet())
    {
      BatchUpdateRequestBuilder builder = new BatchUpdateRequestBuilder(request.getBaseUriTemplate(),
                                                                                    valueType.getType(),
                                                                                    request.getResourceSpec(),
                                                                                    request.getRequestOptions());
      builder.inputs(batch.getValue());
      for (Map.Entry param : request.getQueryParamsObjects().entrySet())
      {
        if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM))
        {
          builder.setParam(param.getKey(), param.getValue());
        }
      }
      for (Map.Entry header : request.getHeaders().entrySet())
      {
        builder.setHeader(header.getKey(), header.getValue());
      }

      RequestContext context = requestContext.clone();
      KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());

      scatterGatherRequests.add(new KVRequestInfo(builder.build(), context));
    }

    return new KVScatterGatherResult(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
  }

  private  MapKeyResult mapKeys(BatchRequest request, Collection ids)
    throws ServiceUnavailableException
  {
    URI serviceUri;
    try
    {
      serviceUri = new URI(D2_URI_PREFIX + request.getServiceName());
    }
    catch (URISyntaxException e)
    {
      throw new IllegalArgumentException(e);
    }

    return _mapper.mapKeysV2(serviceUri, ids);
  }

  /**
   * Helper function to map services to inputs, rather than services to ids.
   * Each input is represented by a Map from keys to {@link RecordTemplate}s.
   *
   * Essentially, instead of calling {@link com.linkedin.d2.balancer.util.MapKeyResult#getMapResult()}
   * to get a map from services to keys, you would instead call this to get a map from services to inputs.
   * You can then use this input to create a new BatchRequest by calling
   * {@link BatchUpdateRequestBuilder#inputs(java.util.Map)} or similar function.
   *
   * @param mapKeyResult {@link MapKeyResult} of mapping U to keys.
   * @param batchRequest the {@link BatchRequest}.
   * @param  the service that will handle each set of inputs; Generally {@link URI}, for a host that will handle the request.
   * @param  the key type.
   * @return a map from U to request input, where request input is a map from keys to {@link RecordTemplate}s.
   */
  private  Map> keyMapToInput(MapKeyResult mapKeyResult, BatchUpdateRequest batchRequest)
  {
    Map updateInput = batchRequest.getUpdateInputMap();

    if (updateInput == null)
    {
      throw new IllegalArgumentException("given BatchRequest must have input data");
    }

    Map> map = mapKeyResult.getMapResult();
    Map> result = new HashMap>(map.size());
    for(Map.Entry> entry : map.entrySet())
    {
      Collection keyList = entry.getValue();
      Map keyRecordMap = new HashMap(keyList.size());
      for(K key : keyList)
      {
        T record = updateInput.get(key);
        if (record == null)
        {
          throw new IllegalArgumentException("given BatchRequest input must have all keys present in mapKeyResult");
        }
        keyRecordMap.put(key, record);
      }
      result.put(entry.getKey(), keyRecordMap);
    }
    return result;
  }

  @SuppressWarnings("deprecation")
  public  KVScatterGatherResult buildRequests(BatchDeleteRequest request, RequestContext requestContext)
    throws ServiceUnavailableException
  {
    Set idObjects = request.getObjectIds();
    Collection ids = new HashSet(idObjects.size());
    for (Object o : idObjects)
    {
      @SuppressWarnings("unchecked")
      K k = (K) o;
      ids.add(k);
    }

    MapKeyResult mapKeyResult = mapKeys(request, ids);
    Map> batches = mapKeyResult.getMapResult();
    Collection> scatterGatherRequests = new ArrayList>(batches.size());

    for (Map.Entry> batch : batches.entrySet())
    {
      TypeSpec value = request.getResourceProperties().getValueType();
      @SuppressWarnings("unchecked")
      Class valueClass = (Class) ((value == null) ? null : value.getType());
      BatchDeleteRequestBuilder builder = new BatchDeleteRequestBuilder(request.getBaseUriTemplate(),
                                                                                    valueClass,
                                                                                    request.getResourceSpec(),
                                                                                    request.getRequestOptions());
      builder.ids(batch.getValue());
      for (Map.Entry param : request.getQueryParamsObjects().entrySet())
      {
        if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM))
        {
          builder.setParam(param.getKey(), param.getValue());
        }
      }
      for (Map.Entry header : request.getHeaders().entrySet())
      {
        builder.setHeader(header.getKey(), header.getValue());
      }

      RequestContext context = requestContext.clone();
      KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());

      BatchRequest> build = builder.build();
      scatterGatherRequests.add(new KVRequestInfo(build, context));
    }

    return new KVScatterGatherResult(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
  }

  /**
   * A convenience function for caller to issue batch request with one call.
   * If finer-grain control is required, users should call buildRequests instead and send requests by themselves
   *
   * @param client - the RestClient to use
   * @param request - the batch get request
   * @param requestContext - the original request context
   * @param callback - callback to be used for each request
   * @throws ServiceUnavailableException
   */
  public void sendRequests(RestClient client, BatchGetRequest request, RequestContext requestContext, Callback>> callback)
    throws ServiceUnavailableException
  {
    ScatterGatherResult scatterGatherResult = buildRequestsV2(request, requestContext);
    for (RequestInfo requestInfo : scatterGatherResult.getRequestInfo())
    {
      client.sendRequest(requestInfo.getRequest(), requestInfo.getRequestContext(), callback);
    }
  }

  /**
   * A convenience function for caller to issue batch request with one call.
   * If finer-grain control is required, users should call buildRequests instead and send requests by themselves
   *
   * @param client - the RestClient to use
   * @param request - the batch get request
   * @param requestContext - the original request context
   * @param callback - callback to be used for each request
   * @throws ServiceUnavailableException
   */
  public  void sendRequests(RestClient client, BatchGetKVRequest request, RequestContext requestContext, Callback>> callback)
      throws ServiceUnavailableException
  {
    KVScatterGatherResult scatterGatherResult = buildRequestsKV(request, requestContext);
    for (KVRequestInfo requestInfo : scatterGatherResult.getRequestInfo())
    {
      client.sendRequest(requestInfo.getRequest(), requestInfo.getRequestContext(), callback);
    }
  }

  public  void sendRequests(RestClient client,
                               BatchUpdateRequest request,
                               RequestContext requestContext,
                               Callback>> callback)
    throws ServiceUnavailableException
  {
    KVScatterGatherResult scatterGatherResult = buildRequests(request, requestContext);
    for(KVRequestInfo requestInfo : scatterGatherResult.getRequestInfo())
    {
      client.sendRequest(requestInfo.getRequest(), requestInfo.getRequestContext(), callback);
    }
  }

  public  void sendRequests(RestClient client,
                               BatchDeleteRequest request,
                               RequestContext requestContext,
                               Callback>> callback)
    throws ServiceUnavailableException
  {
    KVScatterGatherResult scatterGatherResult = buildRequests(request, requestContext);
    for(KVRequestInfo requestInfo : scatterGatherResult.getRequestInfo())
    {
      client.sendRequest(requestInfo.getRequest(), requestInfo.getRequestContext(), callback);
    }
  }

  public static class RequestInfo
  {
    private final BatchRequest> _request;
    private final RequestContext _requestContext;

    public RequestInfo(BatchRequest> request, RequestContext requestContext)
    {
      _request = request;
      _requestContext = requestContext;
    }

    public Request> getRequest()
    {
      return _request;
    }

    public BatchRequest> getBatchRequest()
    {
      return _request;
    }

    public RequestContext getRequestContext()
    {
      return _requestContext;
    }
  }

  public static class ScatterGatherResult
  {
    private final Collection> _requestInfos;
    private final Collection> _unmappedKeys;

    public ScatterGatherResult(Collection> requestInfos, Collection> unmappedKeys)
    {
      _requestInfos = Collections.unmodifiableCollection(requestInfos);
      _unmappedKeys = Collections.unmodifiableCollection(unmappedKeys);
    }

    public Collection> getRequestInfo()
    {
      return _requestInfos;
    }

    public Collection> getUnmappedKeys()
    {
      return _unmappedKeys;
    }
  }

  public static class KVRequestInfo
  {
    private final BatchRequest> _request;
    private final RequestContext _requestContext;

    public KVRequestInfo(BatchRequest> request, RequestContext requestContext)
    {
      _request = request;
      _requestContext = requestContext;
    }

    public BatchRequest> getRequest()
    {
      return _request;
    }

    public RequestContext getRequestContext()
    {
      return _requestContext;
    }
  }

  public static class KVScatterGatherResult
  {
    private final Collection> _requestInfos;
    private final Collection> _unmappedKeys;

    public KVScatterGatherResult(Collection> requestInfos, Collection> unmappedKeys)
    {
      _requestInfos = Collections.unmodifiableCollection(requestInfos);
      _unmappedKeys = Collections.unmodifiableCollection(unmappedKeys);
    }

    public Collection> getRequestInfo()
    {
      return _requestInfos;
    }

    public Collection> getUnmappedKeys()
    {
      return _unmappedKeys;
    }
  }
}