com.hubspot.singularity.resources.AbstractLeaderAwareResource Maven / Gradle / Ivy

Go to download
package com.hubspot.singularity.resources;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Charsets;
import com.ning.http.client.AsyncHttpClient;
import com.ning.http.client.AsyncHttpClient.BoundRequestBuilder;
import com.ning.http.client.Request;
import com.ning.http.client.Response;
import java.io.IOException;
import java.util.Enumeration;
import java.util.concurrent.ExecutionException;
import java.util.function.Supplier;
import javax.servlet.http.HttpServletRequest;
import javax.ws.rs.WebApplicationException;
import org.apache.curator.framework.recipes.leader.LeaderLatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class AbstractLeaderAwareResource {
  private static final Logger LOG = LoggerFactory.getLogger(
    AbstractLeaderAwareResource.class
  );

  protected final AsyncHttpClient httpClient;
  protected final LeaderLatch leaderLatch;
  protected final ObjectMapper objectMapper;

  public AbstractLeaderAwareResource(
    AsyncHttpClient httpClient,
    LeaderLatch leaderLatch,
    ObjectMapper objectMapper
  ) {
    this.httpClient = httpClient;
    this.leaderLatch = leaderLatch;
    this.objectMapper = objectMapper;
  }

  protected boolean isLeader() {
    return leaderLatch.hasLeadership();
  }

  protected  T maybeProxyToLeader(
    HttpServletRequest request,
    Class clazz,
    Q body,
    Supplier runnable
  ) {
    if (leaderLatch.hasLeadership()) {
      return runnable.get();
    }

    String leaderUri;
    try {
      leaderUri = leaderLatch.getLeader().getId();
    } catch (Exception e) {
      throw new RuntimeException("Could not get leader uri to proxy request");
    }

    if (leaderUri.equals(leaderLatch.getId())) {
      LOG.warn(
        "Got own leader id when not the leader! There is likely no leader, will not proxy"
      );
      return runnable.get();
    }

    String url = "http://" + leaderUri + request.getContextPath() + request.getPathInfo();
    LOG.debug("Not the leader, proxying request to {}", url);

    BoundRequestBuilder requestBuilder;
    switch (request.getMethod().toUpperCase()) {
      case "POST":
        requestBuilder = httpClient.preparePost(url);
        break;
      case "PUT":
        requestBuilder = httpClient.preparePut(url);
        break;
      case "DELETE":
        requestBuilder = httpClient.prepareDelete(url);
        break;
      default:
        throw new WebApplicationException(
          String.format("Not meant to proxy request of method %s", request.getMethod()),
          400
        );
    }

    try {
      if (body != null) {
        requestBuilder.setBody(objectMapper.writeValueAsBytes(body));
        LOG.trace("Added body {} to reqeust", body);
      }
    } catch (JsonProcessingException jpe) {
      LOG.error("Could not write body from object {}", body);
      throw new WebApplicationException(jpe, 500);
    }

    copyHeadersAndParams(requestBuilder, request);
    Request httpRequest = requestBuilder.build();

    Response response;
    try {
      LOG.trace("Sending request to leader: {}", httpRequest);
      response = httpClient.executeRequest(httpRequest).get();
    } catch (ExecutionException | InterruptedException e) {
      LOG.error("Could not proxy request {} to leader", e);
      throw new WebApplicationException(e, 500);
    }

    try {
      if (clazz.isAssignableFrom(javax.ws.rs.core.Response.class)) {
        return (T) javax
          .ws.rs.core.Response.status(response.getStatusCode())
          .entity(response.getResponseBody())
          .build();
      }

      if (response.getStatusCode() > 399) {
        throw new WebApplicationException(
          response.getResponseBody(Charsets.UTF_8.toString()),
          response.getStatusCode()
        );
      } else {
        return objectMapper.readValue(response.getResponseBodyAsStream(), clazz);
      }
    } catch (IOException ioe) {
      String message = String.format(
        "Request to leader succeeded with status %s, but could not interpret response",
        response.getStatusCode()
      );
      LOG.error(message, ioe);
      throw new WebApplicationException(message, ioe, 500);
    }
  }

  private void copyHeadersAndParams(
    BoundRequestBuilder requestBuilder,
    HttpServletRequest request
  ) {
    Enumeration headerNames = request.getHeaderNames();
    if (headerNames != null) {
      while (headerNames.hasMoreElements()) {
        String headerName = headerNames.nextElement();
        requestBuilder.addHeader(headerName, request.getHeader(headerName));
        LOG.trace("Copied header {}:{}", headerName, request.getHeader(headerName));
      }
    }
    Enumeration parameterNames = request.getParameterNames();
    if (parameterNames != null) {
      while (parameterNames.hasMoreElements()) {
        String parameterName = parameterNames.nextElement();
        requestBuilder.addQueryParam(parameterName, request.getParameter(parameterName));
        LOG.trace(
          "Copied query param {}={}",
          parameterName,
          request.getParameter(parameterName)
        );
      }
    }
  }

  protected boolean useWebCache(Boolean useWebCache) {
    return useWebCache != null && useWebCache;
  }
}