package org.apache.helix.rest.client;

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.HttpStatus;
import org.apache.http.client.ClientProtocolException;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

class CustomRestClientImpl implements CustomRestClient {
  private static final Logger LOG = LoggerFactory.getLogger(CustomRestClient.class);

  // postfix used to append at the end of base url
  private static final String INSTANCE_HEALTH_STATUS = "/instanceHealthStatus";
  private static final String PARTITION_HEALTH_STATUS = "/partitionHealthStatus";

  private static final String IS_HEALTHY_FIELD = "IS_HEALTHY";
  private static final String PARTITIONS = "partitions";
  private static final String ACCEPT_CONTENT_TYPE = "application/json";
  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  private HttpClient _httpClient;

  private interface JsonConverter {
    Map<String, Boolean> convert(JsonNode jsonNode);
  }

  public CustomRestClientImpl(HttpClient httpClient) {
    _httpClient = httpClient;
  }

  @Override
  public Map<String, Boolean> getInstanceStoppableCheck(String baseUrl,
      Map<String, String> customPayloads) throws IOException {
    // example url: http://<baseUrl>/instanceHealthStatus, assuming the base url already directly
    // queries at the instance
    String url = baseUrl + INSTANCE_HEALTH_STATUS;
    JsonConverter jsonConverter = jsonNode -> {
      Map<String, Boolean> result = new HashMap<>();
      jsonNode.fields().forEachRemaining(kv -> result.put(kv.getKey(), kv.getValue().asBoolean()));
      return result;
    };
    return handleResponse(post(url, customPayloads), jsonConverter);
  }

  @Override
  public Map<String, Boolean> getPartitionStoppableCheck(String baseUrl, List<String> partitions,
      Map<String, String> customPayloads) throws IOException {
    /*
     * example url: http://<baseUrl>/partitionHealthStatus -d {
     * "partitions" : ["p1", "p3", "p9"],
     * "<key>": "<value>",
     * ...
     * }
     */
    String url = baseUrl + PARTITION_HEALTH_STATUS;
    // To avoid ImmutableMap as parameter
    Map<String, String> payLoads = new HashMap<>(customPayloads);
    // Add the entry: "partitions" : ["p1", "p3", "p9"]
    if (partitions != null) {
      payLoads.put(PARTITIONS, partitions.toString());
    }
    JsonConverter jsonConverter = jsonNode -> {
      Map<String, Boolean> result = new HashMap<>();
      jsonNode.fields().forEachRemaining(
          kv -> result.put(kv.getKey(), kv.getValue().get(IS_HEALTHY_FIELD).asBoolean()));
      return result;
    };
    return handleResponse(post(url, payLoads), jsonConverter);
  }

  @VisibleForTesting
  protected JsonNode getJsonObject(HttpResponse httpResponse) throws IOException {
    HttpEntity httpEntity = httpResponse.getEntity();
    String str = EntityUtils.toString(httpEntity);
    LOG.info("Converting Response Content {} to JsonNode", str);
    return OBJECT_MAPPER.readTree(str);
  }

  private Map<String, Boolean> handleResponse(HttpResponse httpResponse,
      JsonConverter jsonConverter) throws IOException {
    int status = httpResponse.getStatusLine().getStatusCode();
    if (status == HttpStatus.SC_OK) {
      LOG.info("Expected HttpResponse statusCode: {}", HttpStatus.SC_OK);
      return jsonConverter.convert(getJsonObject(httpResponse));
    } else {
      throw new ClientProtocolException("Unexpected response status: " + status + ", reason: "
          + httpResponse.getStatusLine().getReasonPhrase());
    }
  }

  @VisibleForTesting
  protected HttpResponse post(String url, Map<String, String> payloads) throws IOException {
    try {
      HttpPost postRequest = new HttpPost(url);
      postRequest.setHeader("Accept", ACCEPT_CONTENT_TYPE);
      StringEntity entity = new StringEntity(OBJECT_MAPPER.writeValueAsString(payloads),
          ContentType.APPLICATION_JSON);
      postRequest.setEntity(entity);
      LOG.info("Executing request: {}, headers: {}, entity: {}", postRequest.getRequestLine(),
          postRequest.getAllHeaders(), postRequest.getEntity());
      return _httpClient.execute(postRequest);
    } catch (IOException e) {
      LOG.error("Failed to perform customized health check. Is participant endpoint {} available?",
          url, e);
      throw e;
    }
  }
}

@@ -1670,16 +1670,20 @@
 stClient
+Impl
 .class);
@@ -4691,16 +4691,138 @@
  else %7B%0A
+      // Ensure entity is fully consumed so stream is closed.%0A      EntityUtils.consumeQuietly(httpResponse.getEntity());%0A
       th
@@ -5098,20 +5098,8 @@
 n %7B%0A
-    try %7B%0A  
     
@@ -5140,16 +5140,26 @@
 t(url);%0A
+    try %7B%0A
       po
@@ -5537,28 +5537,46 @@
 ity());%0A
+%0A
       
-return
+HttpResponse response =
  _httpCl
@@ -5602,16 +5602,235 @@
 quest);%0A
+      int status = response.getStatusLine().getStatusCode();%0A      if (status != HttpStatus.SC_OK) %7B%0A        LOG.warn(%22Received non-200 status code: %7B%7D, payloads: %7B%7D%22, status, payloads);%0A      %7D%0A%0A      return response;%0A
     %7D ca
@@ -5970,16 +5970,126 @@
 rl, e);%0A
+      // Release connection to be reused and avoid connection leakage.%0A      postRequest.releaseConnection();%0A
       th

Fix connection leakage in CustomRestClient (#1472)