From dd28dead2ec461cd6d954f9e270e6bf22463b8cf Mon Sep 17 00:00:00 2001 From: Laszlo Bodor Date: Thu, 27 Oct 2022 20:12:38 +0200 Subject: [PATCH] HIVE-26670: Track every single HTTP request between beeline and hs2 --- .../apache/hadoop/hive/conf/Constants.java | 2 + .../org/apache/hive/jdbc/HiveConnection.java | 28 ++++++++++-- .../hive/jdbc/HttpRequestInterceptorBase.java | 36 +++++++++++++++ .../jdbc/HttpResponseInterceptorBase.java | 44 +++++++++++++++++++ jdbc/src/java/org/apache/hive/jdbc/Utils.java | 2 + .../service/cli/thrift/ThriftHttpServlet.java | 11 ++++- 6 files changed, 119 insertions(+), 4 deletions(-) create mode 100644 jdbc/src/java/org/apache/hive/jdbc/HttpResponseInterceptorBase.java diff --git a/common/src/java/org/apache/hadoop/hive/conf/Constants.java b/common/src/java/org/apache/hadoop/hive/conf/Constants.java index decdc648a6cf..c923c2bcce31 100644 --- a/common/src/java/org/apache/hadoop/hive/conf/Constants.java +++ b/common/src/java/org/apache/hadoop/hive/conf/Constants.java @@ -103,4 +103,6 @@ public class Constants { public static final Pattern COMPACTION_POOLS_PATTERN = Pattern.compile("hive\\.compactor\\.worker\\.(.*)\\.threads"); public static final String HIVE_COMPACTOR_WORKER_POOL = "hive.compactor.worker.pool"; + public static final String HTTP_HEADER_REQUEST_TRACK = "Request-Track"; + public static final String TIME_POSTFIX_REQUEST_TRACK = "_TIME"; } diff --git a/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java b/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java index 4a6fb7c423de..563669970a57 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java +++ b/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java @@ -70,6 +70,7 @@ import java.util.concurrent.Executor; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Stream; +import java.util.function.Supplier; import javax.net.ssl.KeyManagerFactory; import javax.net.ssl.SSLContext; @@ -144,6 +145,7 @@ import org.apache.http.protocol.HttpContext; import org.apache.http.ssl.SSLContexts; import org.apache.http.util.Args; +import org.apache.thrift.TBaseHelper; import org.apache.thrift.TException; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.transport.THttpClient; @@ -182,6 +184,7 @@ public class HiveConnection implements java.sql.Connection { private Subject loggedInSubject; private int maxRetries = 1; private IJdbcBrowserClient browserClient; + private Map additionalHttpHeaders = new HashMap(); /** * Get all direct HiveServer2 URLs from a ZooKeeper based HiveServer2 URL @@ -559,8 +562,7 @@ private CloseableHttpClient getHttpClient(Boolean useSsl) throws SQLException { CookieStore cookieStore = isCookieEnabled ? new BasicCookieStore() : null; HttpClientBuilder httpClientBuilder = null; // Request interceptor for any request pre-processing logic - HttpRequestInterceptor requestInterceptor; - Map additionalHttpHeaders = new HashMap(); + HttpRequestInterceptorBase requestInterceptor; Map customCookies = new HashMap(); // Retrieve the additional HttpHeaders @@ -752,8 +754,12 @@ protected boolean requestIsAborted(final HttpRequest request) { httpClientBuilder .setRedirectStrategy(new HiveJdbcSamlRedirectStrategy(browserClient)); } + + requestInterceptor.setRequestTrackingEnabled(isRequestTrackingEnabled()); + // Add the request interceptor to the client builder - httpClientBuilder.addInterceptorFirst(requestInterceptor); + httpClientBuilder.addInterceptorFirst(requestInterceptor.sessionId(getSessionId())); + httpClientBuilder.addInterceptorLast(new HttpResponseInterceptorBase()); // Add an interceptor to add in an XSRF header httpClientBuilder.addInterceptorLast(new XsrfHttpRequestInterceptor()); @@ -813,6 +819,22 @@ RegistryBuilder. create().register("https", socketFacto return httpClientBuilder.build(); } + private boolean isRequestTrackingEnabled() { + return Boolean.valueOf(sessConfMap.get(JdbcConnectionParams.JDBC_PARAM_REQUEST_TRACK)); + } + + private Supplier getSessionId() { + Supplier sessionId = () -> { + if (sessHandle == null) { + return "NO_SESSION"; + } + StringBuilder b = new StringBuilder(); + TBaseHelper.toString(sessHandle.getSessionId().bufferForGuid(), b); + return b.toString().replaceAll("\\s", ""); + }; + return sessionId; + } + private String getJWT() { String jwtCredential = getJWTStringFromSession(); if (jwtCredential == null || jwtCredential.isEmpty()) { diff --git a/jdbc/src/java/org/apache/hive/jdbc/HttpRequestInterceptorBase.java b/jdbc/src/java/org/apache/hive/jdbc/HttpRequestInterceptorBase.java index eeaa48a5d110..a1c7b58db4f9 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/HttpRequestInterceptorBase.java +++ b/jdbc/src/java/org/apache/hive/jdbc/HttpRequestInterceptorBase.java @@ -19,21 +19,33 @@ package org.apache.hive.jdbc; import java.io.IOException; +import java.util.HashMap; import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Supplier; +import org.apache.hadoop.hive.conf.Constants; import org.apache.http.Header; import org.apache.http.HttpException; import org.apache.http.HttpRequest; import org.apache.http.HttpRequestInterceptor; import org.apache.http.client.CookieStore; import org.apache.http.protocol.HttpContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public abstract class HttpRequestInterceptorBase implements HttpRequestInterceptor { + public static final Logger LOG = LoggerFactory.getLogger(HiveConnection.class.getName()); + CookieStore cookieStore; boolean isCookieEnabled; String cookieName; boolean isSSL; Map additionalHeaders; Map customCookies; + final AtomicLong trackCounter = new AtomicLong(); + private Supplier sessionId = null; + + private boolean requestTrackingEnabled; // Abstract function to add HttpAuth Header protected abstract void addHttpAuthHeader(HttpRequest httpRequest, HttpContext httpContext) @@ -77,6 +89,17 @@ public void process(HttpRequest httpRequest, HttpContext httpContext) if (isCookieEnabled) { httpContext.setAttribute(Utils.HIVE_SERVER2_RETRY_KEY, Utils.HIVE_SERVER2_CONST_FALSE); } + + if (requestTrackingEnabled) { + if (additionalHeaders == null) { + additionalHeaders = new HashMap<>(); + } + String trackHeader = getNewTrackHeader(); + LOG.info("{}:{}", Constants.HTTP_HEADER_REQUEST_TRACK, trackHeader); + additionalHeaders.put(Constants.HTTP_HEADER_REQUEST_TRACK, trackHeader); + httpContext.setAttribute(Constants.HTTP_HEADER_REQUEST_TRACK, trackHeader); + httpContext.setAttribute(trackHeader + Constants.TIME_POSTFIX_REQUEST_TRACK, System.currentTimeMillis()); + } // Insert the additional http headers if (additionalHeaders != null) { for (Map.Entry entry : additionalHeaders.entrySet()) { @@ -102,4 +125,17 @@ public void process(HttpRequest httpRequest, HttpContext httpContext) throw new HttpException(e.getMessage(), e); } } + + protected String getNewTrackHeader() { + return String.format("HIVE_%s_%020d", sessionId.get(), trackCounter.incrementAndGet()); + } + + public HttpRequestInterceptor sessionId(Supplier sessionId) { + this.sessionId = sessionId; + return this; + } + + public void setRequestTrackingEnabled(boolean requestTrackingEnabled) { + this.requestTrackingEnabled = requestTrackingEnabled; + } } diff --git a/jdbc/src/java/org/apache/hive/jdbc/HttpResponseInterceptorBase.java b/jdbc/src/java/org/apache/hive/jdbc/HttpResponseInterceptorBase.java new file mode 100644 index 000000000000..9d1ee8a13c35 --- /dev/null +++ b/jdbc/src/java/org/apache/hive/jdbc/HttpResponseInterceptorBase.java @@ -0,0 +1,44 @@ +/* + * 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. + */ +package org.apache.hive.jdbc; + +import java.io.IOException; + +import org.apache.hadoop.hive.conf.Constants; +import org.apache.http.HttpException; +import org.apache.http.HttpResponse; +import org.apache.http.HttpResponseInterceptor; +import org.apache.http.protocol.HttpContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class HttpResponseInterceptorBase implements HttpResponseInterceptor { + public static final Logger LOG = LoggerFactory.getLogger(HiveConnection.class.getName()); + + @Override + public void process(HttpResponse response, HttpContext context) throws HttpException, IOException { + String trackHeader = (String) context.getAttribute(Constants.HTTP_HEADER_REQUEST_TRACK); + if (trackHeader == null) { + return; + } + long elapsed = System.currentTimeMillis() - (long) context.getAttribute(trackHeader + "_TIME"); + LOG.info("Response to {} in {} ms", trackHeader, elapsed); + context.removeAttribute(Constants.HTTP_HEADER_REQUEST_TRACK); + context.removeAttribute(trackHeader + "_TIME"); + } +} diff --git a/jdbc/src/java/org/apache/hive/jdbc/Utils.java b/jdbc/src/java/org/apache/hive/jdbc/Utils.java index c1be6a52df4f..a855d4e2a5d3 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/Utils.java +++ b/jdbc/src/java/org/apache/hive/jdbc/Utils.java @@ -156,6 +156,8 @@ public static class JdbcConnectionParams { static final String DEFAULT_COOKIE_NAMES_HS2 = "hive.server2.auth"; // The http header prefix for additional headers which have to be appended to the request static final String HTTP_HEADER_PREFIX = "http.header."; + // Request tracking + static final String JDBC_PARAM_REQUEST_TRACK = "requestTrack"; // Set the fetchSize static final String FETCH_SIZE = "fetchSize"; static final String INIT_FILE = "initFile"; diff --git a/service/src/java/org/apache/hive/service/cli/thrift/ThriftHttpServlet.java b/service/src/java/org/apache/hive/service/cli/thrift/ThriftHttpServlet.java index bbb74e0c7ca9..0c192f45ca49 100644 --- a/service/src/java/org/apache/hive/service/cli/thrift/ThriftHttpServlet.java +++ b/service/src/java/org/apache/hive/service/cli/thrift/ThriftHttpServlet.java @@ -24,7 +24,6 @@ import java.nio.charset.StandardCharsets; import java.security.PrivilegedExceptionAction; import java.security.SecureRandom; -import java.text.ParseException; import java.util.Arrays; import java.util.Base64; import java.util.Collections; @@ -45,6 +44,7 @@ import com.google.common.base.Preconditions; import com.google.common.io.ByteStreams; +import org.apache.hadoop.hive.conf.Constants; import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hadoop.hive.conf.HiveConf.ConfVars; import org.apache.hadoop.hive.shims.HadoopShims.KerberosNameShim; @@ -153,6 +153,8 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) String clientIpAddress; boolean requireNewCookie = false; + logTrackingHeaderIfAny(request); + try { if (hiveConf.getBoolean(ConfVars.HIVE_SERVER2_XSRF_FILTER_ENABLED.varname,false)){ boolean continueProcessing = Utils.doXsrfFilter(request,response,null,null); @@ -310,6 +312,13 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) } } + private void logTrackingHeaderIfAny(HttpServletRequest request) { + if (request.getHeader(Constants.HTTP_HEADER_REQUEST_TRACK) != null) { + String requestTrackHeader = request.getHeader(Constants.HTTP_HEADER_REQUEST_TRACK); + LOG.info("{}:{}", Constants.HTTP_HEADER_REQUEST_TRACK, requestTrackHeader); + } + } + private String validateJWT(HttpServletRequest request, HttpServletResponse response) throws HttpAuthenticationException { Preconditions.checkState(jwtValidator != null, "JWT validator should have been set");