diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatch.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatch.java new file mode 100644 index 000000000000..9f60922d6857 --- /dev/null +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatch.java @@ -0,0 +1,24 @@ +/* + * 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.nifi.controller.api.livy; + +public class LivyBatch { + public Integer id; + public String appId; + public String state; + public String[] log; +} diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatchService.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatchService.java new file mode 100644 index 000000000000..1459d5a06467 --- /dev/null +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service-api/src/main/java/org/apache/nifi/controller/api/livy/LivyBatchService.java @@ -0,0 +1,31 @@ +/* + * 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.nifi.controller.api.livy; + +import org.apache.nifi.controller.ControllerService; +import org.apache.nifi.controller.api.livy.exception.SessionManagerException; + +public interface LivyBatchService extends ControllerService { + String APPLICATION_JSON = "application/json"; + String USER = "nifi"; + + LivyBatch getBatchSession(Integer batchid) throws SessionManagerException; + String getBatchStatus(Integer batchid) throws SessionManagerException; + + Integer submitBatch(String codePath, String className, + String codeArgs) throws SessionManagerException; +} diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivyBatchController.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivyBatchController.java new file mode 100644 index 000000000000..647598417c50 --- /dev/null +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivyBatchController.java @@ -0,0 +1,279 @@ +package org.apache.nifi.controller.livy; + +import org.apache.commons.lang3.StringUtils; +import org.apache.nifi.annotation.documentation.CapabilityDescription; +import org.apache.nifi.annotation.documentation.Tags; +import org.apache.nifi.annotation.lifecycle.OnEnabled; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.controller.AbstractControllerService; +import org.apache.nifi.controller.ConfigurationContext; +import org.apache.nifi.controller.ControllerServiceInitializationContext; +import org.apache.nifi.controller.api.livy.LivyBatch; +import org.apache.nifi.controller.api.livy.LivyBatchService; +import org.apache.nifi.controller.api.livy.LivySessionService; +import org.apache.nifi.controller.api.livy.exception.SessionManagerException; +import org.apache.nifi.controller.livy.utilities.LivyHelpers; +import org.apache.nifi.kerberos.KerberosCredentialsService; +import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.ssl.SSLContextService; +import org.codehaus.jackson.map.ObjectMapper; +import org.codehaus.jettison.json.JSONArray; +import org.codehaus.jettison.json.JSONException; +import org.codehaus.jettison.json.JSONObject; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +@Tags({"Livy", "REST", "Spark", "http"}) +@CapabilityDescription("Submits Spark batch jobs over HTTP") +public class LivyBatchController extends AbstractControllerService implements LivyBatchService { + + private volatile String livyUrl; + private volatile String proxyUser; + private volatile String queue; + private volatile String driverMemory; + private volatile Integer driverCores; + private volatile String executorMemory; + private volatile Integer executorCores; + private volatile Integer numExecutors; + private volatile String conf; + + private volatile String jars; + private volatile String files; + private volatile String pyFiles; + private volatile String archives; + private volatile SSLContextService sslContextService; + private volatile int connectTimeout; + private volatile boolean enabled = true; + private volatile KerberosCredentialsService credentialsService; + private volatile String credentialPrincipal; + + private List properties; + + @Override + protected void init(ControllerServiceInitializationContext config) { + final List props = new ArrayList<>(); + props.add(LivyHelpers.LIVY_HOST); + props.add(LivyHelpers.LIVY_PORT); + props.add(LivyHelpers.QUEUE); + props.add(LivyHelpers.SSL_CONTEXT_SERVICE); + props.add(LivyHelpers.CONNECT_TIMEOUT); + props.add(LivyHelpers.JARS); + props.add(LivyHelpers.FILES); + props.add(LivyHelpers.PY_FILES); + props.add(LivyHelpers.ARCHIVES); + props.add(LivyHelpers.KERBEROS_CREDENTIALS_SERVICE); + props.add(LivyHelpers.PROXY_USER); + props.add(LivyHelpers.DRIVER_MEMORY); + props.add(LivyHelpers.DRIVER_CORES); + props.add(LivyHelpers.EXECUTOR_MEMORY); + props.add(LivyHelpers.EXECUTOR_CORES); + props.add(LivyHelpers.EXECUTOR_COUNT); + props.add(LivyHelpers.CONF); + + properties = Collections.unmodifiableList(props); + } + + @OnEnabled + public void onConfigured(final ConfigurationContext context) { + final String livyHost = context.getProperty(LivyHelpers.LIVY_HOST).evaluateAttributeExpressions().getValue(); + final String livyPort = context.getProperty(LivyHelpers.LIVY_PORT).evaluateAttributeExpressions().getValue(); + final String jars = context.getProperty(LivyHelpers.JARS).evaluateAttributeExpressions().getValue(); + final String files = context.getProperty(LivyHelpers.FILES).evaluateAttributeExpressions().getValue(); + final String pyFiles = context.getProperty(LivyHelpers.PY_FILES).evaluateAttributeExpressions().getValue(); + final String archives = context.getProperty(LivyHelpers.ARCHIVES).evaluateAttributeExpressions().getValue(); + + final String proxyUser = context.getProperty(LivyHelpers.PROXY_USER).evaluateAttributeExpressions().getValue(); + final String queue = context.getProperty(LivyHelpers.QUEUE).evaluateAttributeExpressions().getValue(); + + final String driverMemory = context.getProperty(LivyHelpers.DRIVER_MEMORY).evaluateAttributeExpressions().getValue(); + final Integer driverCores = context.getProperty(LivyHelpers.DRIVER_CORES).evaluateAttributeExpressions().asInteger(); + final String executorMemory = context.getProperty(LivyHelpers.EXECUTOR_MEMORY).evaluateAttributeExpressions().getValue(); + final Integer executorCores = context.getProperty(LivyHelpers.EXECUTOR_CORES).evaluateAttributeExpressions().asInteger(); + final Integer numExecutors = context.getProperty(LivyHelpers.EXECUTOR_COUNT).evaluateAttributeExpressions().asInteger(); + final String conf = context.getProperty(LivyHelpers.CONF).evaluateAttributeExpressions().getValue(); + + sslContextService = context.getProperty(LivyHelpers.SSL_CONTEXT_SERVICE).asControllerService(SSLContextService.class); + connectTimeout = Math.toIntExact(context.getProperty(LivyHelpers.CONNECT_TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS)); + credentialsService = context.getProperty(LivyHelpers.KERBEROS_CREDENTIALS_SERVICE).asControllerService(KerberosCredentialsService.class); + + this.livyUrl = "http" + (sslContextService != null ? "s" : "") + "://" + livyHost + ":" + livyPort; + this.jars = jars; + this.files = files; + this.pyFiles = pyFiles; + this.archives = archives; + this.proxyUser = proxyUser; + this.queue = queue; + this.driverMemory = driverMemory; + this.driverCores = driverCores; + this.executorMemory = executorMemory; + this.executorCores = executorCores; + this.numExecutors = numExecutors; + this.conf = conf; + this.enabled = true; + + // Store a copy of the credentialsService principal name for easy string matching + if(credentialsService != null) { + credentialPrincipal = credentialsService.getPrincipal(); + credentialPrincipal = credentialPrincipal.substring(0, credentialPrincipal.indexOf("@")); + } + } + + @Override + protected List getSupportedPropertyDescriptors() { + return properties; + } + + @Override + public LivyBatch getBatchSession(Integer batchid) throws SessionManagerException { + String sessionsUrl = livyUrl + "/batches/" + batchid; + Map headers = new HashMap<>(); + headers.put("Content-Type", APPLICATION_JSON); + headers.put("X-Requested-By", USER); + try { + JSONObject sessionsInfo = LivyHelpers.readJSONFromUrl( + sslContextService, credentialsService, connectTimeout,sessionsUrl, headers); + + LivyBatch livyBatch = new LivyBatch(); + livyBatch.id = batchid; + livyBatch.appId = sessionsInfo.getString("appId"); + livyBatch.state = sessionsInfo.getString("state"); + + // Get log list, which is a list of String + JSONArray logList = sessionsInfo.getJSONArray("log"); + if(logList != null){ + livyBatch.log = new String[logList.length()]; + for(int l=0;l headers = new HashMap<>(); + headers.put("Content-Type", LivySessionService.APPLICATION_JSON); + headers.put("X-Requested-By", LivySessionService.USER); + headers.put("Accept", "application/json"); + + try { + JSONObject batchState = LivyHelpers.readJSONFromUrl(sslContextService, + credentialsService, connectTimeout, livyUrl + "/batches/" + batchid + "/state", + headers); + + return batchState.getString("state"); + } catch(JSONException|IOException e){ + throw new SessionManagerException(e); + } + } + + @Override + public Integer submitBatch(String codePath, String className, + String codeArgs) throws SessionManagerException { + + ComponentLog log = getLogger(); + String batchUrl = livyUrl + "/batches"; + JSONObject output = null; + Map headers = new HashMap<>(); + headers.put("Content-Type", LivySessionService.APPLICATION_JSON); + headers.put("X-Requested-By", LivySessionService.USER); + + log.debug("submitBatch() Submitting Batch Job to Spark via: " + batchUrl); + + StringBuilder payload = new StringBuilder("{\"file\":\"" + codePath + "\""); + appendValue(payload, "className", className); + + try { + if (codeArgs != null) { + payload.append(",\"args\":"); + payload.append(splitCollect(codeArgs)); + } + if (jars != null) { + payload.append(",\"jars\":"); + payload.append(splitCollect(jars)); + } + if (files != null) { + payload.append(",\"files\":"); + payload.append(splitCollect(files)); + } + if (pyFiles != null) { + payload.append(",\"pyFiles\":"); + payload.append(splitCollect(pyFiles)); + } + if (archives != null) { + payload.append(",\"archives\":"); + payload.append(splitCollect(archives)); + } + } catch (IOException e){ + throw new SessionManagerException(e); + } + + appendValue(payload, "proxyUser", proxyUser); + appendValue(payload, "queue", queue); + appendValue(payload, "driverMemory", driverMemory); + appendValue(payload, "driverCores", driverCores); + appendValue(payload, "executorMemory", executorMemory); + appendValue(payload, "executorCores", executorCores); + appendValue(payload, "numExecutors", numExecutors); + appendSubObject(payload, "conf", conf); + + payload.append("}"); + + log.debug("submitBatch() Payload: " + payload.toString()); + + try { + output = LivyHelpers.readJSONObjectFromUrlPOST( + sslContextService, credentialsService, connectTimeout, batchUrl, headers, payload.toString()); + + return output.getInt("id"); + } catch (JSONException|IOException e){ + throw new SessionManagerException(e); + } + } + + final ObjectMapper mapper = new ObjectMapper(); + + private void appendSubObject(StringBuilder payload, String key, String val){ + if (val != null) { + payload.append(",\"" + key + "\": "); + payload.append(val); + } + } + + private void appendValue(StringBuilder payload, String key, String val){ + if (val != null) { + payload.append(",\"" + key + "\": \""); + payload.append(val); + payload.append("\""); + } + } + + private void appendValue(StringBuilder payload, String key, Integer val){ + if (val != null) { + payload.append(",\"" + key + "\": "); + payload.append(val); + } + } + + private String splitCollect(String arg) throws IOException { + List filesArray = Arrays.stream(arg.split(",")) + .filter(StringUtils::isNotBlank) + .map(String::trim).collect(Collectors.toList()); + String filesJsonArray = mapper.writeValueAsString(filesArray); + + return filesJsonArray; + } +} diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivySessionController.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivySessionController.java index 77b146a1c2d1..880b5d3ad258 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivySessionController.java +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/LivySessionController.java @@ -16,20 +16,9 @@ */ package org.apache.nifi.controller.livy; -import java.io.BufferedReader; -import java.io.FileInputStream; import java.io.IOException; -import java.io.InputStream; -import java.io.InputStreamReader; import java.net.ConnectException; import java.net.SocketTimeoutException; -import java.nio.charset.StandardCharsets; -import java.security.KeyManagementException; -import java.security.KeyStore; -import java.security.KeyStoreException; -import java.security.NoSuchAlgorithmException; -import java.security.UnrecoverableKeyException; -import java.security.cert.CertificateException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -40,24 +29,8 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; -import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; -import org.apache.http.HttpEntity; -import org.apache.http.HttpResponse; -import org.apache.http.HttpStatus; -import org.apache.http.auth.AuthSchemeProvider; -import org.apache.http.auth.AuthScope; -import org.apache.http.client.CredentialsProvider; -import org.apache.http.client.HttpClient; -import org.apache.http.client.config.AuthSchemes; -import org.apache.http.client.config.RequestConfig; -import org.apache.http.client.methods.HttpGet; -import org.apache.http.client.methods.HttpPost; -import org.apache.http.config.Lookup; -import org.apache.http.config.RegistryBuilder; -import org.apache.http.entity.StringEntity; -import org.apache.http.impl.client.BasicCredentialsProvider; -import org.apache.http.impl.client.HttpClientBuilder; +import org.apache.http.impl.client.CloseableHttpClient; import org.apache.nifi.annotation.documentation.CapabilityDescription; import org.apache.nifi.annotation.documentation.Tags; import org.apache.nifi.annotation.lifecycle.OnDisabled; @@ -67,8 +40,7 @@ import org.apache.nifi.controller.ConfigurationContext; import org.apache.nifi.controller.ControllerServiceInitializationContext; import org.apache.nifi.controller.api.livy.exception.SessionManagerException; -import org.apache.nifi.hadoop.KerberosKeytabCredentials; -import org.apache.nifi.hadoop.KerberosKeytabSPNegoAuthSchemeProvider; +import org.apache.nifi.controller.livy.utilities.LivyHelpers; import org.apache.nifi.kerberos.KerberosCredentialsService; import org.apache.nifi.logging.ComponentLog; import org.apache.nifi.processor.util.StandardValidators; @@ -80,51 +52,40 @@ import org.apache.nifi.controller.api.livy.LivySessionService; import org.apache.nifi.expression.ExpressionLanguageScope; -import javax.net.ssl.KeyManagerFactory; -import javax.net.ssl.SSLContext; -import javax.net.ssl.TrustManagerFactory; - @Tags({"Livy", "REST", "Spark", "http"}) @CapabilityDescription("Manages pool of Spark sessions over HTTP") public class LivySessionController extends AbstractControllerService implements LivySessionService { - - public static final PropertyDescriptor LIVY_HOST = new PropertyDescriptor.Builder() - .name("livy-cs-livy-host") - .displayName("Livy Host") - .description("The hostname (or IP address) of the Livy server.") - .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) - .required(true) - .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) - .build(); - - public static final PropertyDescriptor LIVY_PORT = new PropertyDescriptor.Builder() - .name("livy-cs-livy-port") - .displayName("Livy Port") - .description("The port number for the Livy server.") - .required(true) - .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) - .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) - .defaultValue("8998") - .build(); - public static final PropertyDescriptor SESSION_POOL_SIZE = new PropertyDescriptor.Builder() .name("livy-cs-session-pool-size") - .displayName("Session Pool Size") - .description("Number of sessions to keep open") + .displayName("Min Session Pool Size") + .description("Minimum number of sessions to keep open. If Elastic Session Pool is enabled, new sessions " + + "will be started automatically if no idle Livy sessions are available, up to the " + + "Max Session Pool Size.") .required(true) .defaultValue("2") .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) .build(); - public static final PropertyDescriptor SESSION_TYPE = new PropertyDescriptor.Builder() - .name("livy-cs-session-kind") - .displayName("Session Type") - .description("The type of Spark session to start (spark, pyspark, pyspark3, sparkr, e.g.)") + public static final PropertyDescriptor ELASTIC_SESSION_POOL_SIZE = new PropertyDescriptor.Builder() + .name("livy-cs-elastic-session-pool-size") + .displayName("Elastic Session Pool") + .description("When enabled, if no 'idle' Livy sessions are available, a new session will be created, " + + "up to the Max Session Pool Size.") .required(true) - .allowableValues("spark", "pyspark", "pyspark3", "sparkr") - .defaultValue("spark") - .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .allowableValues("true", "false") + .defaultValue("true") + .build(); + + public static final PropertyDescriptor MAX_SESSION_POOL_SIZE = new PropertyDescriptor.Builder() + .name("livy-cs-max-session-pool-size") + .displayName("Max Session Pool Size") + .description("Maximum number of sessions to keep open. Only used when Elastic Session Pool is enabled." + + "A value of 0 means there is no limit on the number of Livy sessions.") + .required(false) + .defaultValue("0") + .addValidator(StandardValidators.NON_NEGATIVE_INTEGER_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) .build(); public static final PropertyDescriptor SESSION_MGR_STATUS_INTERVAL = new PropertyDescriptor.Builder() @@ -137,60 +98,65 @@ public class LivySessionController extends AbstractControllerService implements .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) .build(); - public static final PropertyDescriptor JARS = new PropertyDescriptor.Builder() - .name("livy-cs-session-jars") - .displayName("Session JARs") - .description("JARs to be used in the Spark session.") - .required(false) - .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) - .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) - .build(); - - public static final PropertyDescriptor FILES = new PropertyDescriptor.Builder() - .name("livy-cs-session-files") - .displayName("Session Files") - .description("Files to be used in the Spark session.") + public static final PropertyDescriptor SESSION_NAME = new PropertyDescriptor.Builder() + .name("livy-cs-session-name") + .displayName("Session Name") + .description("Session Name is used to differentiate Livy Sessions, and ensure the controller service uses " + + "only sessions it is supposed to manage. If empty, defaults to the Controller Service UUID.") .required(false) - .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) - .defaultValue(null) .build(); - public static final PropertyDescriptor SSL_CONTEXT_SERVICE = new PropertyDescriptor.Builder() - .name("SSL Context Service") - .description("The SSL Context Service used to provide client certificate information for TLS/SSL (https) connections.") - .required(false) - .identifiesControllerService(SSLContextService.class) + public static final PropertyDescriptor CLEANUP_SESSIONS = new PropertyDescriptor.Builder() + .name("livy-cs-session-cleanup") + .displayName("Cleanup Sessions") + .description("When the Livy Controller Service shuts down, should created Livy sessions be deleted?") + .allowableValues("true", "false") + .defaultValue("false") + .required(true) .build(); - public static final PropertyDescriptor CONNECT_TIMEOUT = new PropertyDescriptor.Builder() - .name("Connection Timeout") - .description("Max wait time for connection to remote service.") + public static final PropertyDescriptor KEEP_SESSIONS_ALIVE = new PropertyDescriptor.Builder() + .name("livy-cs-session-keep-alive") + .displayName("Keep Busy Sessions Alive") + .description("If a session is busy, should NiFi prevent Livy from terminating that session when the Livy " + + "idle timeout is reached?") + .allowableValues("true", "false") + .defaultValue("true") .required(true) - .defaultValue("5 secs") - .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) .build(); - static final PropertyDescriptor KERBEROS_CREDENTIALS_SERVICE = new PropertyDescriptor.Builder() - .name("kerberos-credentials-service") - .displayName("Kerberos Credentials Service") - .description("Specifies the Kerberos Credentials Controller Service that should be used for authenticating with Kerberos") - .identifiesControllerService(KerberosCredentialsService.class) - .required(false) - .build(); - private volatile String livyUrl; private volatile int sessionPoolSize; + private volatile boolean elasticSessionPool; + private volatile int maxSessionPoolSize; private volatile String controllerKind; + private volatile String proxyUser; + private volatile String queue; + private volatile String driverMemory; + private volatile Integer driverCores; + private volatile String executorMemory; + private volatile Integer executorCores; + private volatile Integer numExecutors; + private volatile String sessionName; + private volatile String conf; + private volatile boolean cleanupSessions; + private volatile boolean keepSessionsAlive; + + private volatile String jars; private volatile String files; + private volatile String pyFiles; + private volatile String archives; private volatile Map sessions = new ConcurrentHashMap<>(); private volatile SSLContextService sslContextService; - private volatile SSLContext sslContext; private volatile int connectTimeout; private volatile Thread livySessionManagerThread = null; private volatile boolean enabled = true; private volatile KerberosCredentialsService credentialsService; + private volatile String credentialPrincipal; + private volatile String appropriateProxy; private volatile SessionManagerException sessionManagerException; private List properties; @@ -198,16 +164,31 @@ public class LivySessionController extends AbstractControllerService implements @Override protected void init(ControllerServiceInitializationContext config) { final List props = new ArrayList<>(); - props.add(LIVY_HOST); - props.add(LIVY_PORT); + props.add(LivyHelpers.LIVY_HOST); + props.add(LivyHelpers.LIVY_PORT); props.add(SESSION_POOL_SIZE); - props.add(SESSION_TYPE); + props.add(ELASTIC_SESSION_POOL_SIZE); + props.add(MAX_SESSION_POOL_SIZE); + props.add(CLEANUP_SESSIONS); + props.add(KEEP_SESSIONS_ALIVE); + props.add(LivyHelpers.SESSION_TYPE); + props.add(SESSION_NAME); + props.add(LivyHelpers.QUEUE); props.add(SESSION_MGR_STATUS_INTERVAL); - props.add(SSL_CONTEXT_SERVICE); - props.add(CONNECT_TIMEOUT); - props.add(JARS); - props.add(FILES); - props.add(KERBEROS_CREDENTIALS_SERVICE); + props.add(LivyHelpers.SSL_CONTEXT_SERVICE); + props.add(LivyHelpers.CONNECT_TIMEOUT); + props.add(LivyHelpers.JARS); + props.add(LivyHelpers.FILES); + props.add(LivyHelpers.PY_FILES); + props.add(LivyHelpers.ARCHIVES); + props.add(LivyHelpers.KERBEROS_CREDENTIALS_SERVICE); + props.add(LivyHelpers.PROXY_USER); + props.add(LivyHelpers.DRIVER_MEMORY); + props.add(LivyHelpers.DRIVER_CORES); + props.add(LivyHelpers.EXECUTOR_MEMORY); + props.add(LivyHelpers.EXECUTOR_CORES); + props.add(LivyHelpers.EXECUTOR_COUNT); + props.add(LivyHelpers.CONF); properties = Collections.unmodifiableList(props); } @@ -219,25 +200,64 @@ protected List getSupportedPropertyDescriptors() { @OnEnabled public void onConfigured(final ConfigurationContext context) { - final String livyHost = context.getProperty(LIVY_HOST).evaluateAttributeExpressions().getValue(); - final String livyPort = context.getProperty(LIVY_PORT).evaluateAttributeExpressions().getValue(); - final String sessionPoolSize = context.getProperty(SESSION_POOL_SIZE).evaluateAttributeExpressions().getValue(); - final String sessionKind = context.getProperty(SESSION_TYPE).getValue(); + final String livyHost = context.getProperty(LivyHelpers.LIVY_HOST).evaluateAttributeExpressions().getValue(); + final String livyPort = context.getProperty(LivyHelpers.LIVY_PORT).evaluateAttributeExpressions().getValue(); + final Integer sessionPoolSize = context.getProperty(SESSION_POOL_SIZE).evaluateAttributeExpressions().asInteger(); + final boolean elasticSessionPool = context.getProperty(ELASTIC_SESSION_POOL_SIZE).asBoolean(); + final Integer maxSessionPoolSize = context.getProperty(MAX_SESSION_POOL_SIZE).evaluateAttributeExpressions().asInteger(); + final String sessionKind = context.getProperty(LivyHelpers.SESSION_TYPE).getValue(); final long sessionManagerStatusInterval = context.getProperty(SESSION_MGR_STATUS_INTERVAL).evaluateAttributeExpressions().asTimePeriod(TimeUnit.MILLISECONDS); - final String jars = context.getProperty(JARS).evaluateAttributeExpressions().getValue(); - final String files = context.getProperty(FILES).evaluateAttributeExpressions().getValue(); - sslContextService = context.getProperty(SSL_CONTEXT_SERVICE).asControllerService(SSLContextService.class); - sslContext = sslContextService == null ? null : sslContextService.createSSLContext(SSLContextService.ClientAuth.NONE); - connectTimeout = Math.toIntExact(context.getProperty(CONNECT_TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS)); - credentialsService = context.getProperty(KERBEROS_CREDENTIALS_SERVICE).asControllerService(KerberosCredentialsService.class); + final String jars = context.getProperty(LivyHelpers.JARS).evaluateAttributeExpressions().getValue(); + final String files = context.getProperty(LivyHelpers.FILES).evaluateAttributeExpressions().getValue(); + final String pyFiles = context.getProperty(LivyHelpers.PY_FILES).evaluateAttributeExpressions().getValue(); + final String archives = context.getProperty(LivyHelpers.ARCHIVES).evaluateAttributeExpressions().getValue(); + + final String proxyUser = context.getProperty(LivyHelpers.PROXY_USER).evaluateAttributeExpressions().getValue(); + final String queue = context.getProperty(LivyHelpers.QUEUE).evaluateAttributeExpressions().getValue(); + + final String driverMemory = context.getProperty(LivyHelpers.DRIVER_MEMORY).evaluateAttributeExpressions().getValue(); + final Integer driverCores = context.getProperty(LivyHelpers.DRIVER_CORES).evaluateAttributeExpressions().asInteger(); + final String executorMemory = context.getProperty(LivyHelpers.EXECUTOR_MEMORY).evaluateAttributeExpressions().getValue(); + final Integer executorCores = context.getProperty(LivyHelpers.EXECUTOR_CORES).evaluateAttributeExpressions().asInteger(); + final Integer numExecutors = context.getProperty(LivyHelpers.EXECUTOR_COUNT).evaluateAttributeExpressions().asInteger(); + final String sessionName = context.getProperty(SESSION_NAME).evaluateAttributeExpressions().getValue(); + final String conf = context.getProperty(LivyHelpers.CONF).evaluateAttributeExpressions().getValue(); + final boolean cleanupSessions = context.getProperty(CLEANUP_SESSIONS).asBoolean(); + final boolean keepSessionsAlive = context.getProperty(KEEP_SESSIONS_ALIVE).asBoolean(); + + sslContextService = context.getProperty(LivyHelpers.SSL_CONTEXT_SERVICE).asControllerService(SSLContextService.class); + connectTimeout = Math.toIntExact(context.getProperty(LivyHelpers.CONNECT_TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS)); + credentialsService = context.getProperty(LivyHelpers.KERBEROS_CREDENTIALS_SERVICE).asControllerService(KerberosCredentialsService.class); this.livyUrl = "http" + (sslContextService != null ? "s" : "") + "://" + livyHost + ":" + livyPort; this.controllerKind = sessionKind; this.jars = jars; this.files = files; - this.sessionPoolSize = Integer.valueOf(sessionPoolSize); + this.pyFiles = pyFiles; + this.archives = archives; + this.sessionPoolSize = sessionPoolSize; + this.maxSessionPoolSize = maxSessionPoolSize; + this.elasticSessionPool = elasticSessionPool; + this.proxyUser = proxyUser; + this.queue = queue; + this.driverMemory = driverMemory; + this.driverCores = driverCores; + this.executorMemory = executorMemory; + this.executorCores = executorCores; + this.numExecutors = numExecutors; + this.sessionName = StringUtils.isEmpty(sessionName)?this.getIdentifier():sessionName; + this.conf = conf; + this.cleanupSessions = cleanupSessions; + this.keepSessionsAlive = keepSessionsAlive; this.enabled = true; + // Store a copy of the credentialsService principal name for easy string matching + if(credentialsService != null) { + credentialPrincipal = credentialsService.getPrincipal(); + credentialPrincipal = credentialPrincipal.substring(0, credentialPrincipal.indexOf("@")); + appropriateProxy = StringUtils.isEmpty(proxyUser)?credentialPrincipal:proxyUser; + } + livySessionManagerThread = new Thread(() -> { while (enabled) { try { @@ -266,6 +286,18 @@ public void shutdown() { enabled = false; livySessionManagerThread.interrupt(); livySessionManagerThread.join(); + + // If we need to cleanup opened sessions + if(this.cleanupSessions){ + for(Integer s : sessions.keySet()) { + // Don't re-throw session cleanup errors, just log + try { + JSONObject deleteMessage = this.deleteSession(s); + } catch (IOException e) { + log.warn("Livy Session Manager failed to cleanup session #" + s, e); + } + } + } } catch (InterruptedException e) { Thread.currentThread().interrupt(); log.error("Livy Session Manager Thread interrupted"); @@ -296,45 +328,16 @@ public Map getSession() throws SessionManagerException { } @Override - public HttpClient getConnection() throws IOException, SessionManagerException { + public CloseableHttpClient getConnection() throws IOException, SessionManagerException { checkSessionManagerException(); - return openConnection(); - } - - private HttpClient openConnection() throws IOException { - HttpClientBuilder httpClientBuilder = HttpClientBuilder.create(); - - if (sslContextService != null) { - try { - SSLContext sslContext = getSslSocketFactory(sslContextService); - httpClientBuilder.setSSLContext(sslContext); - } catch (KeyStoreException | CertificateException | NoSuchAlgorithmException | UnrecoverableKeyException | KeyManagementException e) { - throw new IOException(e); - } - } - - if (credentialsService != null) { - CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); - credentialsProvider.setCredentials(new AuthScope(null, -1, null), - new KerberosKeytabCredentials(credentialsService.getPrincipal(), credentialsService.getKeytab())); - httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider); - Lookup authSchemeRegistry = RegistryBuilder. create() - .register(AuthSchemes.SPNEGO, new KerberosKeytabSPNegoAuthSchemeProvider()).build(); - httpClientBuilder.setDefaultAuthSchemeRegistry(authSchemeRegistry); - } - - RequestConfig.Builder requestConfigBuilder = RequestConfig.custom(); - requestConfigBuilder.setConnectTimeout(connectTimeout); - requestConfigBuilder.setConnectionRequestTimeout(connectTimeout); - requestConfigBuilder.setSocketTimeout(connectTimeout); - httpClientBuilder.setDefaultRequestConfig(requestConfigBuilder.build()); - - return httpClientBuilder.build(); + return LivyHelpers.openConnection(sslContextService, credentialsService, connectTimeout); } private void manageSessions() throws InterruptedException, IOException { int idleSessions = 0; + int startingSessions = 0; + Map idleSessionInfo = new HashMap<>(); JSONObject newSessionInfo; Map sessionsInfo; ComponentLog log = getLogger(); @@ -348,24 +351,58 @@ private void manageSessions() throws InterruptedException, IOException { for (Integer sessionId : new ArrayList<>(sessions.keySet())) { JSONObject currentSession = sessions.get(sessionId); log.debug("manageSessions() Updating current session: " + currentSession); + if (sessionsInfo.containsKey(sessionId)) { String state = currentSession.getString("state"); String sessionKind = currentSession.getString("kind"); - log.debug("manageSessions() controller kind: {}, session kind: {}, session state: {}", - new Object[]{controllerKind, sessionKind, state}); - if (state.equalsIgnoreCase("idle") && sessionKind.equalsIgnoreCase(controllerKind)) { + + String sessionOwner = currentSession.has("owner")?currentSession.getString("owner"):""; + String sessionProxy = currentSession.has("proxyUser")?currentSession.getString("proxyUser"):""; + + log.debug("manageSessions() session ID: {}, session owner: {}, session proxy: {}, controller kind: {}, session kind: {}, session state: {}", + new Object[]{sessionId,sessionOwner, sessionProxy, controllerKind, sessionKind, state}); + + if (!sessionKind.equalsIgnoreCase(controllerKind)) { + // Prune sessions of kind != controllerKind, and where Owner is not contained in principal name + sessions.remove(sessionId); + //Remove session from session list source of truth snapshot since it has been dealt with + sessionsInfo.remove(sessionId); + continue; + } + + // If security is enabled, do owner/proxy checks + if (credentialsService != null && (!credentialPrincipal.equals(sessionOwner) + || !appropriateProxy.equals(sessionProxy))) { + + log.debug("manageSessions() Dropping session, not owned or proxied by this account. session ID: {}", + new Object[]{sessionId, sessionOwner, sessionProxy, controllerKind, sessionKind, state}); + // where Owner or Proxy is not contained in principal name + sessions.remove(sessionId); + //Remove session from session list source of truth snapshot since it has been dealt with + sessionsInfo.remove(sessionId); + continue; + } + + // Status filters + if (state.equalsIgnoreCase("idle")) { // Keep track of how many sessions are in an idle state and thus available idleSessions++; + idleSessionInfo.put(sessionId, sessionsInfo.get(sessionId)); sessions.put(sessionId, sessionsInfo.get(sessionId)); // Remove session from session list source of truth snapshot since it has been dealt with sessionsInfo.remove(sessionId); - } else if ((state.equalsIgnoreCase("busy") || state.equalsIgnoreCase("starting")) && sessionKind.equalsIgnoreCase(controllerKind)) { + } else if ((state.equalsIgnoreCase("busy") || state.equalsIgnoreCase("starting"))) { + // Track starting instance count + if(state.equalsIgnoreCase("starting")){ + startingSessions++; + } + // Update status of existing sessions sessions.put(sessionId, sessionsInfo.get(sessionId)); // Remove session from session list source of truth snapshot since it has been dealt with sessionsInfo.remove(sessionId); } else { - // Prune sessions of kind != controllerKind and whose state is: + // Prune sessions whose state is: // not_started, shutting_down, error, dead, success (successfully stopped) sessions.remove(sessionId); //Remove session from session list source of truth snapshot since it has been dealt with @@ -389,9 +426,22 @@ private void manageSessions() throws InterruptedException, IOException { log.debug("manageSessions() Registered new session: " + newSessionInfo); } } else { - // Open one new session if there are no idle sessions - if (idleSessions == 0) { - log.debug("manageSessions() There are " + numSessions + " sessions in the pool but none of them are idle sessions, creating..."); + // If we exceeded our session pool size, look for `idle` sessions we can shut down + // Two scenarios: we have no elastic pool sizing, in which case definitely look for candidates + // Or, we do have elastic pool sizing, in which case we need to make sure we are above our max + // pool size + if(idleSessions > 0 && numSessions > sessionPoolSize && (!elasticSessionPool || numSessions > maxSessionPoolSize)) { + int sessionID = idleSessionInfo.keySet().stream().findFirst().get(); + log.debug("manageSessions() There are " + numSessions + " sessions in the pool, " + + "this exceeds the maximum number of allowed sessions, shutting down idle Livy session " + sessionID + "..."); + deleteSession(sessionID); + } + + // If Elastic Session Pool sizing is enabled + // Open one new session if there are no idle or starting sessions + if (elasticSessionPool && numSessions < maxSessionPoolSize && (idleSessions + startingSessions) == 0) { + log.debug("manageSessions() There are " + numSessions + " sessions in the pool, " + + "none of them are idle sessions and Elastic Session Pool Sizing is enabled, creating..."); newSessionInfo = openSession(); sessions.put(newSessionInfo.getInt("id"), newSessionInfo); log.debug("manageSessions() Registered new session: " + newSessionInfo); @@ -405,6 +455,19 @@ private void manageSessions() throws InterruptedException, IOException { log.debug("manageSessions() Registered new session: " + newSessionInfo); } } + + // Sessions that are running may need to be "connected" to remotely to reset their timeouts. + // Otherwise Livy will kill them based on bug LIVY-547. + if(keepSessionsAlive){ + for(Map.Entry me : sessionsInfo.entrySet()) { + // Check if session is busy before updating it's activity date + String state = me.getValue().getString("state"); + if(state.equalsIgnoreCase("busy")){ + getSessionConnect(me.getKey()); + log.debug("manageSessions() Connected to session to keep it alive: " + me.getKey()); + } + } + } } } catch (ConnectException | SocketTimeoutException ce) { log.error("Timeout connecting to Livy service to retrieve sessions", ce); @@ -422,7 +485,8 @@ private Map listSessions() throws IOException { headers.put("Content-Type", APPLICATION_JSON); headers.put("X-Requested-By", USER); try { - sessionsInfo = readJSONFromUrl(sessionsUrl, headers); + sessionsInfo = LivyHelpers.readJSONFromUrl( + sslContextService, credentialsService, connectTimeout,sessionsUrl, headers); numSessions = sessionsInfo.getJSONArray("sessions").length(); for (int i = 0; i < numSessions; i++) { int currentSessionId = sessionsInfo.getJSONArray("sessions").getJSONObject(i).getInt("id"); @@ -443,7 +507,24 @@ private JSONObject getSessionInfo(int sessionId) throws IOException { headers.put("Content-Type", APPLICATION_JSON); headers.put("X-Requested-By", USER); try { - sessionInfo = readJSONFromUrl(sessionUrl, headers); + sessionInfo = LivyHelpers.readJSONFromUrl( + sslContextService, credentialsService, connectTimeout, sessionUrl, headers); + } catch (JSONException e) { + throw new IOException(e); + } + + return sessionInfo; + } + + private JSONObject getSessionConnect(int sessionId) throws IOException { + String sessionUrl = livyUrl + "/sessions/" + sessionId + "/connect"; + JSONObject sessionInfo; + Map headers = new HashMap<>(); + headers.put("Content-Type", APPLICATION_JSON); + headers.put("X-Requested-By", USER); + try { + sessionInfo = LivyHelpers.readJSONObjectFromUrlPOST( + sslContextService, credentialsService, connectTimeout, sessionUrl, headers, null); } catch (JSONException e) { throw new IOException(e); } @@ -475,6 +556,63 @@ private JSONObject openSession() throws IOException, JSONException, InterruptedE payload.append(",\"files\":"); payload.append(filesJsonArray); } + if (pyFiles != null) { + List filesArray = Arrays.stream(pyFiles.split(",")) + .filter(StringUtils::isNotBlank) + .map(String::trim).collect(Collectors.toList()); + String filesJsonArray = mapper.writeValueAsString(filesArray); + payload.append(",\"pyFiles\":"); + payload.append(filesJsonArray); + } + if (archives != null) { + List filesArray = Arrays.stream(archives.split(",")) + .filter(StringUtils::isNotBlank) + .map(String::trim).collect(Collectors.toList()); + String filesJsonArray = mapper.writeValueAsString(filesArray); + payload.append(",\"archives\":"); + payload.append(filesJsonArray); + } + if (proxyUser != null) { + payload.append(",\"proxyUser\": \""); + payload.append(proxyUser); + payload.append("\""); + } + if (queue != null) { + payload.append(",\"queue\": \""); + payload.append(queue); + payload.append("\""); + } + if (driverMemory != null) { + payload.append(",\"driverMemory\": \""); + payload.append(driverMemory); + payload.append("\""); + } + if (driverCores != null) { + payload.append(",\"driverCores\": "); + payload.append(driverCores); + } + if (executorMemory != null) { + payload.append(",\"executorMemory\": \""); + payload.append(executorMemory); + payload.append("\""); + } + if (executorCores != null) { + payload.append(",\"executorCores\": "); + payload.append(executorCores); + } + if (numExecutors != null) { + payload.append(",\"numExecutors\": "); + payload.append(numExecutors); + } + if (sessionName != null) { + payload.append(",\"name\": \""); + payload.append(sessionName); + payload.append("\""); + } + if (conf != null) { + payload.append(",\"conf\": "); + payload.append(conf); + } payload.append("}"); log.debug("openSession() Session Payload: " + payload.toString()); @@ -482,7 +620,8 @@ private JSONObject openSession() throws IOException, JSONException, InterruptedE headers.put("Content-Type", APPLICATION_JSON); headers.put("X-Requested-By", USER); - newSessionInfo = readJSONObjectFromUrlPOST(sessionsUrl, headers, payload.toString()); + newSessionInfo = LivyHelpers.readJSONObjectFromUrlPOST( + sslContextService, credentialsService, connectTimeout, sessionsUrl, headers, payload.toString()); Thread.sleep(1000); while (newSessionInfo.getString("state").equalsIgnoreCase("starting")) { log.debug("openSession() Waiting for session to start..."); @@ -494,73 +633,20 @@ private JSONObject openSession() throws IOException, JSONException, InterruptedE return newSessionInfo; } - private JSONObject readJSONObjectFromUrlPOST(String urlString, Map headers, String payload) throws IOException, JSONException { - HttpClient httpClient = openConnection(); - - HttpPost request = new HttpPost(urlString); - for (Map.Entry entry : headers.entrySet()) { - request.addHeader(entry.getKey(), entry.getValue()); - } - HttpEntity httpEntity = new StringEntity(payload); - request.setEntity(httpEntity); - HttpResponse response = httpClient.execute(request); - - if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK && response.getStatusLine().getStatusCode() != HttpStatus.SC_CREATED) { - throw new RuntimeException("Failed : HTTP error code : " + response.getStatusLine().getStatusCode() + " : " + response.getStatusLine().getReasonPhrase()); - } - - InputStream content = response.getEntity().getContent(); - return readAllIntoJSONObject(content); - } - - private JSONObject readJSONFromUrl(String urlString, Map headers) throws IOException, JSONException { - HttpClient httpClient = openConnection(); - - HttpGet request = new HttpGet(urlString); - for (Map.Entry entry : headers.entrySet()) { - request.addHeader(entry.getKey(), entry.getValue()); - } - HttpResponse response = httpClient.execute(request); - - InputStream content = response.getEntity().getContent(); - return readAllIntoJSONObject(content); - } - - private JSONObject readAllIntoJSONObject(InputStream content) throws IOException, JSONException { - BufferedReader rd = new BufferedReader(new InputStreamReader(content, StandardCharsets.UTF_8)); - String jsonText = IOUtils.toString(rd); - return new JSONObject(jsonText); - } - - private SSLContext getSslSocketFactory(SSLContextService sslService) - throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException, UnrecoverableKeyException, KeyManagementException { - final String keystoreLocation = sslService.getKeyStoreFile(); - final String keystorePass = sslService.getKeyStorePassword(); - final String keystoreType = sslService.getKeyStoreType(); - - // prepare the keystore - final KeyStore keyStore = KeyStore.getInstance(keystoreType); - - try (FileInputStream keyStoreStream = new FileInputStream(keystoreLocation)) { - keyStore.load(keyStoreStream, keystorePass.toCharArray()); + private JSONObject deleteSession(int sessionId) throws IOException { + String sessionUrl = livyUrl + "/sessions/" + sessionId; + JSONObject sessionInfo; + Map headers = new HashMap<>(); + headers.put("Content-Type", APPLICATION_JSON); + headers.put("X-Requested-By", USER); + try { + sessionInfo = LivyHelpers.readJSONFromUrlDELETE( + sslContextService, credentialsService, connectTimeout, sessionUrl, headers); + } catch (JSONException e) { + throw new IOException(e); } - final KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); - keyManagerFactory.init(keyStore, keystorePass.toCharArray()); - - // load truststore - final String truststoreLocation = sslService.getTrustStoreFile(); - final String truststorePass = sslService.getTrustStorePassword(); - final String truststoreType = sslService.getTrustStoreType(); - - KeyStore truststore = KeyStore.getInstance(truststoreType); - final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance("X509"); - truststore.load(new FileInputStream(truststoreLocation), truststorePass.toCharArray()); - trustManagerFactory.init(truststore); - - sslContext.init(keyManagerFactory.getKeyManagers(), trustManagerFactory.getTrustManagers(), null); - - return sslContext; + return sessionInfo; } private void checkSessionManagerException() throws SessionManagerException { diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/utilities/LivyHelpers.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/utilities/LivyHelpers.java new file mode 100644 index 000000000000..17bf23d4425d --- /dev/null +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/java/org/apache/nifi/controller/livy/utilities/LivyHelpers.java @@ -0,0 +1,357 @@ +package org.apache.nifi.controller.livy.utilities; + +import org.apache.commons.io.IOUtils; +import org.apache.http.HttpEntity; +import org.apache.http.HttpResponse; +import org.apache.http.HttpStatus; +import org.apache.http.auth.AuthSchemeProvider; +import org.apache.http.auth.AuthScope; +import org.apache.http.client.CredentialsProvider; +import org.apache.http.client.config.AuthSchemes; +import org.apache.http.client.config.RequestConfig; +import org.apache.http.client.methods.HttpDelete; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.client.methods.HttpPost; +import org.apache.http.config.Lookup; +import org.apache.http.config.RegistryBuilder; +import org.apache.http.entity.StringEntity; +import org.apache.http.impl.client.BasicCredentialsProvider; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClientBuilder; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.hadoop.KerberosKeytabCredentials; +import org.apache.nifi.hadoop.KerberosKeytabSPNegoAuthSchemeProvider; +import org.apache.nifi.kerberos.KerberosCredentialsService; +import org.apache.nifi.processor.util.StandardValidators; +import org.apache.nifi.ssl.SSLContextService; +import org.codehaus.jettison.json.JSONException; +import org.codehaus.jettison.json.JSONObject; + +import javax.net.ssl.KeyManagerFactory; +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManagerFactory; +import java.io.BufferedReader; +import java.io.FileInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.security.KeyManagementException; +import java.security.KeyStore; +import java.security.KeyStoreException; +import java.security.NoSuchAlgorithmException; +import java.security.UnrecoverableKeyException; +import java.security.cert.CertificateException; +import java.util.Map; + +public class LivyHelpers { + public static final PropertyDescriptor LIVY_HOST = new PropertyDescriptor.Builder() + .name("livy-cs-livy-host") + .displayName("Livy Host") + .description("The hostname (or IP address) of the Livy server.") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .required(true) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor LIVY_PORT = new PropertyDescriptor.Builder() + .name("livy-cs-livy-port") + .displayName("Livy Port") + .description("The port number for the Livy server.") + .required(true) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .defaultValue("8998") + .build(); + + public static final PropertyDescriptor SESSION_TYPE = new PropertyDescriptor.Builder() + .name("livy-cs-session-kind") + .displayName("Session Type") + .description("The type of Spark session to start (spark, pyspark, pyspark3, sparkr, e.g.)") + .required(true) + .allowableValues("spark", "pyspark", "pyspark3", "sparkr") + .defaultValue("spark") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + public static final PropertyDescriptor PROXY_USER = new PropertyDescriptor.Builder() + .name("livy-cs-session-proxyuser") + .displayName("Proxy User") + .description("User to impersonate when starting the session") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor QUEUE = new PropertyDescriptor.Builder() + .name("livy-cs-session-queue") + .displayName("YARN Queue") + .description("The name of the YARN queue to which submitted") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor DRIVER_MEMORY = new PropertyDescriptor.Builder() + .name("livy-cs-session-driver-memory") + .displayName("Driver Memory") + .description("Amount of memory to use for the driver process") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor DRIVER_CORES = new PropertyDescriptor.Builder() + .name("livy-cs-session-driver-cores") + .displayName("Driver Cores") + .description("Number of cores to use for the driver process") + .required(false) + .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor EXECUTOR_MEMORY = new PropertyDescriptor.Builder() + .name("livy-cs-session-executor-memory") + .displayName("Executor Memory") + .description("Amount of memory to use per executor process") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor EXECUTOR_CORES = new PropertyDescriptor.Builder() + .name("livy-cs-session-executor-cores") + .displayName("Executor Cores") + .description("Number of cores to use for each executor") + .required(false) + .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor EXECUTOR_COUNT = new PropertyDescriptor.Builder() + .name("livy-cs-session-executor-count") + .displayName("Executor Count") + .description("Number of executors to launch for this session") + .required(false) + .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor CONF = new PropertyDescriptor.Builder() + .name("livy-cs-session-conf") + .displayName("Conf") + .description("Spark configuration properties, in JSON format. Ex: {\"spark.driver.maxResultSize\": \"0\", \"spark.submit.deployMode\": \"cluster\"}") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor JARS = new PropertyDescriptor.Builder() + .name("livy-cs-session-jars") + .displayName("Session JARs") + .description("JARs to be used in the Spark session.") + .required(false) + .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .build(); + + public static final PropertyDescriptor FILES = new PropertyDescriptor.Builder() + .name("livy-cs-session-files") + .displayName("Session Files") + .description("Files to be used in the Spark session.") + .required(false) + .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .defaultValue(null) + .build(); + + public static final PropertyDescriptor PY_FILES = new PropertyDescriptor.Builder() + .name("livy-cs-session-py-files") + .displayName("Session Python Files") + .description("Python files to be used in this session") + .required(false) + .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .defaultValue(null) + .build(); + + public static final PropertyDescriptor ARCHIVES = new PropertyDescriptor.Builder() + .name("livy-cs-session-archives") + .displayName("Session Archives") + .description("Archives to be used in this session") + .required(false) + .addValidator(StandardValidators.createListValidator(true, true, StandardValidators.FILE_EXISTS_VALIDATOR)) + .expressionLanguageSupported(ExpressionLanguageScope.VARIABLE_REGISTRY) + .defaultValue(null) + .build(); + + public static final PropertyDescriptor CHARSET = new PropertyDescriptor.Builder() + .name("exec-spark-iactive-charset") + .displayName("Character Set") + .description("The character set encoding for the incoming flow file.") + .required(true) + .defaultValue("UTF-8") + .addValidator(StandardValidators.CHARACTER_SET_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + + public static final PropertyDescriptor SSL_CONTEXT_SERVICE = new PropertyDescriptor.Builder() + .name("SSL Context Service") + .description("The SSL Context Service used to provide client certificate information for TLS/SSL (https) connections.") + .required(false) + .identifiesControllerService(SSLContextService.class) + .build(); + + public static final PropertyDescriptor CONNECT_TIMEOUT = new PropertyDescriptor.Builder() + .name("Connection Timeout") + .description("Max wait time for connection to remote service.") + .required(true) + .defaultValue("5 secs") + .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) + .build(); + + public static final PropertyDescriptor KERBEROS_CREDENTIALS_SERVICE = new PropertyDescriptor.Builder() + .name("kerberos-credentials-service") + .displayName("Kerberos Credentials Service") + .description("Specifies the Kerberos Credentials Controller Service that should be used for authenticating with Kerberos") + .identifiesControllerService(KerberosCredentialsService.class) + .required(false) + .build(); + + public static CloseableHttpClient openConnection(SSLContextService sslContextService, + KerberosCredentialsService credentialsService, + int connectTimeout) throws IOException { + HttpClientBuilder httpClientBuilder = HttpClientBuilder.create(); + + if (sslContextService != null) { + try { + SSLContext sslContext = getSslSocketFactory(sslContextService); + httpClientBuilder.setSSLContext(sslContext); + } catch (KeyStoreException | CertificateException | NoSuchAlgorithmException | UnrecoverableKeyException | KeyManagementException e) { + throw new IOException(e); + } + } + + if (credentialsService != null) { + CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); + credentialsProvider.setCredentials(new AuthScope(null, -1, null), + new KerberosKeytabCredentials(credentialsService.getPrincipal(), credentialsService.getKeytab())); + httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider); + Lookup authSchemeRegistry = RegistryBuilder. create() + .register(AuthSchemes.SPNEGO, new KerberosKeytabSPNegoAuthSchemeProvider()).build(); + httpClientBuilder.setDefaultAuthSchemeRegistry(authSchemeRegistry); + } + + RequestConfig.Builder requestConfigBuilder = RequestConfig.custom(); + requestConfigBuilder.setConnectTimeout(connectTimeout); + requestConfigBuilder.setConnectionRequestTimeout(connectTimeout); + requestConfigBuilder.setSocketTimeout(connectTimeout); + httpClientBuilder.setDefaultRequestConfig(requestConfigBuilder.build()); + + return httpClientBuilder.build(); + } + + public static JSONObject readJSONObjectFromUrlPOST( + SSLContextService sslContextService, + KerberosCredentialsService credentialsService, + int connectTimeout, + String urlString, Map headers, String payload) throws IOException, JSONException { + + try(CloseableHttpClient httpClient = openConnection(sslContextService, credentialsService, connectTimeout)){ + HttpPost request = new HttpPost(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + + if(payload != null) { + HttpEntity httpEntity = new StringEntity(payload); + request.setEntity(httpEntity); + } + + HttpResponse response = httpClient.execute(request); + + if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK && response.getStatusLine().getStatusCode() != HttpStatus.SC_CREATED) { + throw new RuntimeException("Failed : HTTP error code : " + response.getStatusLine().getStatusCode() + " : " + response.getStatusLine().getReasonPhrase()); + } + + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); + } + } + + public static JSONObject readJSONFromUrl( + SSLContextService sslContextService, + KerberosCredentialsService credentialsService, + int connectTimeout, + String urlString, Map headers) throws IOException, JSONException { + + try(CloseableHttpClient httpClient = openConnection(sslContextService, credentialsService, connectTimeout)) { + HttpGet request = new HttpGet(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + HttpResponse response = httpClient.execute(request); + + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); + } + } + + public static JSONObject readJSONFromUrlDELETE( + SSLContextService sslContextService, + KerberosCredentialsService credentialsService, + int connectTimeout, + String urlString, Map headers) throws IOException, JSONException { + + try(CloseableHttpClient httpClient = openConnection(sslContextService, credentialsService, connectTimeout)) { + HttpDelete request = new HttpDelete(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + HttpResponse response = httpClient.execute(request); + + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); + } + } + + public static JSONObject readAllIntoJSONObject(InputStream content) throws IOException, JSONException { + BufferedReader rd = new BufferedReader(new InputStreamReader(content, StandardCharsets.UTF_8)); + String jsonText = IOUtils.toString(rd); + return new JSONObject(jsonText); + } + + public static SSLContext getSslSocketFactory(SSLContextService sslService) + throws IOException, KeyStoreException, CertificateException, NoSuchAlgorithmException, UnrecoverableKeyException, KeyManagementException { + + SSLContext sslContext = sslService.createSSLContext(SSLContextService.ClientAuth.NONE); + final String keystoreLocation = sslService.getKeyStoreFile(); + final String keystorePass = sslService.getKeyStorePassword(); + final String keystoreType = sslService.getKeyStoreType(); + + // prepare the keystore + final KeyStore keyStore = KeyStore.getInstance(keystoreType); + + try (FileInputStream keyStoreStream = new FileInputStream(keystoreLocation)) { + keyStore.load(keyStoreStream, keystorePass.toCharArray()); + } + + final KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + keyManagerFactory.init(keyStore, keystorePass.toCharArray()); + + // load truststore + final String truststoreLocation = sslService.getTrustStoreFile(); + final String truststorePass = sslService.getTrustStorePassword(); + final String truststoreType = sslService.getTrustStoreType(); + + KeyStore truststore = KeyStore.getInstance(truststoreType); + final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance("X509"); + truststore.load(new FileInputStream(truststoreLocation), truststorePass.toCharArray()); + trustManagerFactory.init(truststore); + + sslContext.init(keyManagerFactory.getKeyManagers(), trustManagerFactory.getTrustManagers(), null); + + return sslContext; + } +} diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService index 8b4b0468a745..a96db00ac39f 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-controller-service/src/main/resources/META-INF/services/org.apache.nifi.controller.ControllerService @@ -13,4 +13,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -org.apache.nifi.controller.livy.LivySessionController \ No newline at end of file +org.apache.nifi.controller.livy.LivySessionController +org.apache.nifi.controller.livy.LivyBatchController \ No newline at end of file diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkBatch.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkBatch.java new file mode 100644 index 000000000000..2b1427699be8 --- /dev/null +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkBatch.java @@ -0,0 +1,188 @@ +/* + * 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.nifi.processors.livy; + +import org.apache.nifi.annotation.behavior.InputRequirement; +import org.apache.nifi.annotation.documentation.CapabilityDescription; +import org.apache.nifi.annotation.documentation.Tags; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.controller.api.livy.LivyBatch; +import org.apache.nifi.controller.api.livy.LivyBatchService; +import org.apache.nifi.controller.api.livy.exception.SessionManagerException; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.processor.AbstractProcessor; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.ProcessorInitializationContext; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.util.StandardValidators; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +@InputRequirement(InputRequirement.Requirement.INPUT_ALLOWED) +@Tags({"spark", "livy", "http", "execute"}) +@CapabilityDescription("Execute Spark Code over a Livy-managed HTTP session as a Spark Batch job.") +public class ExecuteSparkBatch extends AbstractProcessor { + public static final PropertyDescriptor LIVY_CONTROLLER_SERVICE = new PropertyDescriptor.Builder() + .name("exec-spark-batch-livy-controller-service") + .displayName("Livy Controller Service") + .description("The controller service to use for submitting Livy-managed batch jobs.") + .required(true) + .identifiesControllerService(LivyBatchService.class) + .build(); + + public static final PropertyDescriptor CODE_PATH = new PropertyDescriptor.Builder() + .name("exec-spark-code-file-path") + .displayName("Code File") + .description("File containing the application to execute. " + + "For HDFS locations use the format hdfs:///path/to/examples.jar") + .required(true) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + + public static final PropertyDescriptor CLASS_NAME = new PropertyDescriptor.Builder() + .name("exec-spark-code-class-name") + .displayName("Class Name") + .description("For JVM applications, the name of the class to execute.") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + + public static final PropertyDescriptor CODE_ARGS = new PropertyDescriptor.Builder() + .name("exec-spark-code-args") + .displayName("Args") + .description("The comma delimited list of arguments to pass to the file to be executed.") + .required(false) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + + public static final PropertyDescriptor STATUS_CHECK_INTERVAL = new PropertyDescriptor.Builder() + .name("exec-spark-iactive-status-check-interval") + .displayName("Status Check Interval") + .description("The amount of time to wait between checking the status of an operation.") + .required(true) + .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .defaultValue("1 sec") + .build(); + + public static final Relationship REL_SUCCESS = new Relationship.Builder() + .name("success") + .description("FlowFiles that are successfully processed are sent to this relationship") + .build(); + + public static final Relationship REL_FAILURE = new Relationship.Builder() + .name("failure") + .description("FlowFiles are routed to this relationship when they cannot be parsed or failed to run") + .build(); + + private volatile List properties; + private volatile Set relationships; + + @Override + public void init(final ProcessorInitializationContext context) { + List properties = new ArrayList<>(); + properties.add(LIVY_CONTROLLER_SERVICE); + properties.add(CODE_PATH); + properties.add(CLASS_NAME); + properties.add(CODE_ARGS); + properties.add(STATUS_CHECK_INTERVAL); + this.properties = Collections.unmodifiableList(properties); + + Set relationships = new HashSet<>(); + relationships.add(REL_SUCCESS); + relationships.add(REL_FAILURE); + this.relationships = Collections.unmodifiableSet(relationships); + } + + @Override + public Set getRelationships() { + return relationships; + } + + @Override + protected List getSupportedPropertyDescriptors() { + return properties; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + + final ComponentLog log = getLogger(); + final LivyBatchService livyBatchService = context.getProperty(LIVY_CONTROLLER_SERVICE).asControllerService(LivyBatchService.class); + + final String codePath = context.getProperty(CODE_PATH).evaluateAttributeExpressions(flowFile).getValue(); + final String className = context.getProperty(CLASS_NAME).evaluateAttributeExpressions(flowFile).getValue(); + final String codeArgs = context.getProperty(CODE_ARGS).evaluateAttributeExpressions(flowFile).getValue(); + + final long statusCheckInterval = context.getProperty(STATUS_CHECK_INTERVAL).evaluateAttributeExpressions(flowFile).asTimePeriod(TimeUnit.MILLISECONDS); + + try { + final Integer batchId = livyBatchService.submitBatch(codePath, className, codeArgs); + String jobState = livyBatchService.getBatchStatus(batchId); + + log.debug("executeSparkBatch() Job status is: " + jobState + ". Waiting for job to complete..."); + + flowFile = session.putAttribute(flowFile, "livy.batchid", Integer.toString(batchId)); + + while (!jobState.equalsIgnoreCase("error") + && !jobState.equalsIgnoreCase("dead") + && !jobState.equalsIgnoreCase("killed") + && !jobState.equalsIgnoreCase("success")) { + log.debug("executeSparkBatch() Job status is: " + jobState + ". Waiting for job to complete..."); + Thread.sleep(statusCheckInterval); + jobState = livyBatchService.getBatchStatus(batchId); + } + + // Livy Batch mode is like Spark-Submit, so unlike Interactive mode there is no output to write to the FlowFile + // We'll write the log output instead + LivyBatch livyBatch = livyBatchService.getBatchSession(batchId); + flowFile = session.putAttribute(flowFile, "livy.jobstate", livyBatch.state); + + if(livyBatch.log != null){ + flowFile = session.write(flowFile, out -> out.write(String.join("\n", livyBatch.log).getBytes())); + } + + // Job has ended, handle error/cancel states + if(jobState.equalsIgnoreCase("error") + || jobState.equalsIgnoreCase("dead") + || jobState.equalsIgnoreCase("killed")) { + session.transfer(flowFile, REL_FAILURE); + } else { + session.transfer(flowFile, REL_SUCCESS); + } + } catch (SessionManagerException|InterruptedException e){ + log.error("Job processing error: " + e); + session.transfer(flowFile, REL_FAILURE); + } + } +} diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkInteractive.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkInteractive.java index b9edcacfe4af..43cfb0e6f1b4 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkInteractive.java +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/java/org/apache/nifi/processors/livy/ExecuteSparkInteractive.java @@ -38,10 +38,11 @@ import org.apache.http.HttpEntity; import org.apache.http.HttpResponse; import org.apache.http.HttpStatus; -import org.apache.http.client.HttpClient; +import org.apache.http.client.methods.HttpDelete; import org.apache.http.client.methods.HttpGet; import org.apache.http.client.methods.HttpPost; import org.apache.http.entity.StringEntity; +import org.apache.http.impl.client.CloseableHttpClient; import org.apache.nifi.annotation.behavior.InputRequirement; import org.apache.nifi.annotation.documentation.CapabilityDescription; import org.apache.nifi.annotation.documentation.Tags; @@ -104,6 +105,17 @@ public class ExecuteSparkInteractive extends AbstractProcessor { .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) .build(); + static final PropertyDescriptor DROP_SESSIONS_WITH_ERROR = new PropertyDescriptor.Builder() + .name("exec-spark-iactive-session-error-drop") + .displayName("Drop Sessions On Error") + .description("If a statement ends in an error, should the session be dropped and replaced with a new one?" + + "Some code will cause a session to be left in an unstable state if it errors, and it's better to " + + "start over.") + .allowableValues("true", "false") + .defaultValue("false") + .required(true) + .build(); + public static final PropertyDescriptor STATUS_CHECK_INTERVAL = new PropertyDescriptor.Builder() .name("exec-spark-iactive-status-check-interval") .displayName("Status Check Interval") @@ -139,6 +151,7 @@ public void init(final ProcessorInitializationContext context) { properties.add(CODE); properties.add(CHARSET); properties.add(STATUS_CHECK_INTERVAL); + properties.add(DROP_SESSIONS_WITH_ERROR); this.properties = Collections.unmodifiableList(properties); Set relationships = new HashSet<>(); @@ -185,9 +198,9 @@ public void onTrigger(ProcessContext context, final ProcessSession session) thro } final long statusCheckInterval = context.getProperty(STATUS_CHECK_INTERVAL).evaluateAttributeExpressions(flowFile).asTimePeriod(TimeUnit.MILLISECONDS); Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue()); - - String sessionId = livyController.get("sessionId"); - String livyUrl = livyController.get("livyUrl"); + final boolean dropSessionOnError = context.getProperty(DROP_SESSIONS_WITH_ERROR).asBoolean(); + final String sessionId = livyController.get("sessionId"); + final String livyUrl = livyController.get("livyUrl"); String code = context.getProperty(CODE).evaluateAttributeExpressions(flowFile).getValue(); if (StringUtils.isEmpty(code)) { try (InputStream inputStream = session.read(flowFile)) { @@ -202,28 +215,53 @@ public void onTrigger(ProcessContext context, final ProcessSession session) thro } code = StringEscapeUtils.escapeJson(code); - String payload = "{\"code\":\"" + code + "\"}"; + final String payload = "{\"code\":\"" + code + "\"}"; try { final JSONObject result = submitAndHandleJob(livyUrl, livySessionService, sessionId, payload, statusCheckInterval); - log.debug("ExecuteSparkInteractive Result of Job Submit: " + result); + if (result == null) { + log.debug("ExecuteSparkInteractive No Job Result received: Session ID" + sessionId); session.transfer(flowFile, REL_FAILURE); + return; + } + + log.debug("ExecuteSparkInteractive Result of Job Submit: " + result); + + final int statementid = result.getInt("id"); + final String jobState = result.getString("state"); + + flowFile = session.putAttribute(flowFile, "livy.sessionid", sessionId); + flowFile = session.putAttribute(flowFile, "livy.statementid", Integer.toString(statementid)); + + // Route error/cancel states to failure + if (jobState.equalsIgnoreCase("error") + || jobState.equalsIgnoreCase("cancelled") + || jobState.equalsIgnoreCase("cancelling")) { + throw new IOException(jobState); } else { + final JSONObject output = result.getJSONObject("output"); + try { - final JSONObject output = result.getJSONObject("data"); - flowFile = session.write(flowFile, out -> out.write(output.toString().getBytes(charset))); + final JSONObject data = output.getJSONObject("data"); + flowFile = session.write(flowFile, out -> out.write(data.toString().getBytes(charset))); flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON); session.transfer(flowFile, REL_SUCCESS); } catch (JSONException je) { // The result doesn't contain the data, just send the output object as the flow file content to failure (after penalizing) log.error("Spark Session returned an error, sending the output JSON object as the flow file content to failure (after penalizing)"); - flowFile = session.write(flowFile, out -> out.write(result.toString().getBytes(charset))); + flowFile = session.write(flowFile, out -> out.write(output.toString().getBytes(charset))); flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON); flowFile = session.penalize(flowFile); + + // If we are supposed to drop sessions that caused failures, run it here + if(dropSessionOnError) { + deleteSession(livyUrl, livySessionService, sessionId); + } + session.transfer(flowFile, REL_FAILURE); } } - } catch (IOException | SessionManagerException e) { + } catch (IOException | SessionManagerException | JSONException e) { log.error("Failure processing flowfile {} due to {}, penalizing and routing to failure", new Object[]{flowFile, e.getMessage()}, e); flowFile = session.penalize(flowFile); session.transfer(flowFile, REL_FAILURE); @@ -253,7 +291,7 @@ private JSONObject submitAndHandleJob(String livyUrl, LivySessionService livySes Thread.sleep(statusCheckInterval); if (jobState.equalsIgnoreCase("available")) { log.debug("submitAndHandleJob() Job status is: " + jobState + ". returning output..."); - output = jobInfo.getJSONObject("output"); + output = jobInfo; } else if (jobState.equalsIgnoreCase("running") || jobState.equalsIgnoreCase("waiting")) { while (!jobState.equalsIgnoreCase("available")) { log.debug("submitAndHandleJob() Job status is: " + jobState + ". Waiting for job to complete..."); @@ -261,12 +299,14 @@ private JSONObject submitAndHandleJob(String livyUrl, LivySessionService livySes jobInfo = readJSONObjectFromUrl(statementUrl, livySessionService, headers); jobState = jobInfo.getString("state"); } - output = jobInfo.getJSONObject("output"); + output = jobInfo; } else if (jobState.equalsIgnoreCase("error") || jobState.equalsIgnoreCase("cancelled") || jobState.equalsIgnoreCase("cancelling")) { log.debug("Job status is: " + jobState + ". Job did not complete due to error or has been cancelled. Check SparkUI for details."); - throw new IOException(jobState); + + // Return the job Info even if it ended in an error or cancel state + output = jobInfo; } } catch (JSONException | InterruptedException e) { throw new IOException(e); @@ -274,37 +314,66 @@ private JSONObject submitAndHandleJob(String livyUrl, LivySessionService livySes return output; } - private JSONObject readJSONObjectFromUrlPOST(String urlString, LivySessionService livySessionService, Map headers, String payload) - throws IOException, JSONException, SessionManagerException { - HttpClient httpClient = livySessionService.getConnection(); + private JSONObject deleteSession(String livyUrl, LivySessionService livySessionService, String sessionId) throws IOException, SessionManagerException { + String sessionUrl = livyUrl + "/sessions/" + sessionId; + JSONObject sessionInfo; + Map headers = new HashMap<>(); + headers.put("Content-Type", LivySessionService.APPLICATION_JSON); + headers.put("X-Requested-By", LivySessionService.USER); - HttpPost request = new HttpPost(urlString); - for (Map.Entry entry : headers.entrySet()) { - request.addHeader(entry.getKey(), entry.getValue()); + try { + sessionInfo = readJSONFromUrlDELETE(sessionUrl, livySessionService, headers); + } catch (JSONException e) { + throw new IOException(e); } - HttpEntity httpEntity = new StringEntity(payload); - request.setEntity(httpEntity); - HttpResponse response = httpClient.execute(request); - if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK && response.getStatusLine().getStatusCode() != HttpStatus.SC_CREATED) { - throw new RuntimeException("Failed : HTTP error code : " + response.getStatusLine().getStatusCode() + " : " + response.getStatusLine().getReasonPhrase()); + return sessionInfo; + } + + private JSONObject readJSONFromUrlDELETE(String urlString, LivySessionService livySessionService, Map headers) throws IOException, JSONException, SessionManagerException { + try(CloseableHttpClient httpClient = (CloseableHttpClient)livySessionService.getConnection()) { + HttpDelete request = new HttpDelete(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + HttpResponse response = httpClient.execute(request); + + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); } + } + + private JSONObject readJSONObjectFromUrlPOST(String urlString, LivySessionService livySessionService, Map headers, String payload) + throws IOException, JSONException, SessionManagerException { + try(CloseableHttpClient httpClient = (CloseableHttpClient)livySessionService.getConnection()) { + HttpPost request = new HttpPost(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + HttpEntity httpEntity = new StringEntity(payload); + request.setEntity(httpEntity); + HttpResponse response = httpClient.execute(request); - InputStream content = response.getEntity().getContent(); - return readAllIntoJSONObject(content); + if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK && response.getStatusLine().getStatusCode() != HttpStatus.SC_CREATED) { + throw new RuntimeException("Failed : HTTP error code : " + response.getStatusLine().getStatusCode() + " : " + response.getStatusLine().getReasonPhrase()); + } + + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); + } } private JSONObject readJSONObjectFromUrl(String urlString, LivySessionService livySessionService, Map headers) throws IOException, JSONException, SessionManagerException { - HttpClient httpClient = livySessionService.getConnection(); + try(CloseableHttpClient httpClient = (CloseableHttpClient)livySessionService.getConnection()) { + HttpGet request = new HttpGet(urlString); + for (Map.Entry entry : headers.entrySet()) { + request.addHeader(entry.getKey(), entry.getValue()); + } + HttpResponse response = httpClient.execute(request); - HttpGet request = new HttpGet(urlString); - for (Map.Entry entry : headers.entrySet()) { - request.addHeader(entry.getKey(), entry.getValue()); + InputStream content = response.getEntity().getContent(); + return readAllIntoJSONObject(content); } - HttpResponse response = httpClient.execute(request); - - InputStream content = response.getEntity().getContent(); - return readAllIntoJSONObject(content); } private JSONObject readAllIntoJSONObject(InputStream content) throws IOException, JSONException { diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor index 8de81e451d63..a659fdbd8131 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor @@ -13,4 +13,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -org.apache.nifi.processors.livy.ExecuteSparkInteractive \ No newline at end of file +org.apache.nifi.processors.livy.ExecuteSparkInteractive +org.apache.nifi.processors.livy.ExecuteSparkBatch \ No newline at end of file diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/ExecuteSparkInteractiveTestBase.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/ExecuteSparkInteractiveTestBase.java index 18fbdebe75b3..1c5341fae1c8 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/ExecuteSparkInteractiveTestBase.java +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/ExecuteSparkInteractiveTestBase.java @@ -19,6 +19,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.commons.io.IOUtils; +import org.apache.nifi.controller.livy.LivySessionController; import org.apache.nifi.web.util.TestServer; import org.apache.nifi.util.MockFlowFile; import org.apache.nifi.util.TestRunner; @@ -56,16 +57,16 @@ public void handle(String target, Request baseRequest, HttpServletRequest reques } else if ("/sessions/1/statements/7".equalsIgnoreCase(target)) { switch (session1Requests) { case 0: - responseBody = "{\"state\": \"waiting\"}"; + responseBody = "{\"id\": 7, \"state\": \"waiting\"}"; break; case 1: - responseBody = "{\"state\": \"running\"}"; + responseBody = "{\"id\": 7, \"state\": \"running\"}"; break; case 2: - responseBody = "{\"state\": \"available\", \"output\": {\"data\": {\"text/plain\": \"Hello world\"}}}"; + responseBody = "{\"id\": 7, \"state\": \"available\", \"output\": {\"data\": {\"text/plain\": \"Hello world\"}}}"; break; default: - responseBody = "{\"state\": \"error\"}"; + responseBody = "{\"id\": 7, \"state\": \"error\"}"; break; } session1Requests++; @@ -81,6 +82,8 @@ public void handle(String target, Request baseRequest, HttpServletRequest reques responseContentType = "application/json"; if ("/sessions".equalsIgnoreCase(target)) { responseBody = "{\"id\": 1, \"kind\": \"spark\", \"state\": \"idle\"}"; + } else if ("/sessions/1/connect".equalsIgnoreCase(target)) { + responseBody = "{\"id\": 1, \"kind\": \"spark\", \"state\": \"idle\"}"; } else if ("/sessions/1/statements".equalsIgnoreCase(target)) { responseBody = "{\"id\": 7}"; } @@ -89,6 +92,23 @@ public void handle(String target, Request baseRequest, HttpServletRequest reques responseContentType = "text/plain"; responseBody = "Bad request"; } + } else if ("DELETE".equalsIgnoreCase(request.getMethod())) { + String requestBody = IOUtils.toString(request.getReader()); + try { + // validate JSON payload + new ObjectMapper().readTree(requestBody); + + responseStatus = 200; + responseBody = "{}"; + responseContentType = "application/json"; + if ("/sessions/1".equalsIgnoreCase(target)) { + responseBody = "{\"msg\":\"deleted\"}"; + } + } catch (JsonProcessingException e) { + responseStatus = 400; + responseContentType = "text/plain"; + responseBody = "Bad request"; + } } response.setStatus(responseStatus); @@ -104,8 +124,12 @@ public void handle(String target, Request baseRequest, HttpServletRequest reques } TestRunner runner; + LivySessionController livyControllerService; void testCode(TestServer server, String code) throws Exception { + runner.enableControllerService(livyControllerService); + runner.setProperty(ExecuteSparkInteractive.LIVY_CONTROLLER_SERVICE, "livyCS"); + server.addHandler(new LivyAPIHandler()); runner.enqueue(code); diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractive.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractive.java index 992e2e5c6ce7..423de1093cf1 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractive.java +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractive.java @@ -17,6 +17,7 @@ package org.apache.nifi.processors.livy; import org.apache.nifi.controller.livy.LivySessionController; +import org.apache.nifi.controller.livy.utilities.LivyHelpers; import org.apache.nifi.util.TestRunners; import org.apache.nifi.web.util.TestServer; import org.junit.After; @@ -52,12 +53,10 @@ public static void afterClass() throws Exception { @Before public void before() throws Exception { runner = TestRunners.newTestRunner(ExecuteSparkInteractive.class); - LivySessionController livyControllerService = new LivySessionController(); + livyControllerService = new LivySessionController(); runner.addControllerService("livyCS", livyControllerService); - runner.setProperty(livyControllerService, LivySessionController.LIVY_HOST, url.substring(url.indexOf("://") + 3, url.lastIndexOf(":"))); - runner.setProperty(livyControllerService, LivySessionController.LIVY_PORT, url.substring(url.lastIndexOf(":") + 1)); - runner.enableControllerService(livyControllerService); - runner.setProperty(ExecuteSparkInteractive.LIVY_CONTROLLER_SERVICE, "livyCS"); + runner.setProperty(livyControllerService, LivyHelpers.LIVY_HOST, url.substring(url.indexOf("://") + 3, url.lastIndexOf(":"))); + runner.setProperty(livyControllerService, LivyHelpers.LIVY_PORT, url.substring(url.lastIndexOf(":") + 1)); server.clearHandlers(); } @@ -80,4 +79,14 @@ public void testSparkSession() throws Exception { public void testSparkSessionWithSpecialChars() throws Exception { testCode(server, "print \"/'?!<>[]{}()$&*=%;.|_-\\\""); } + + @Test + public void testSparkSessionCleanup() throws Exception { + runner.setProperty(livyControllerService, LivySessionController.CLEANUP_SESSIONS, "true"); + + testCode(server, "print \"hello world\""); + + // Disable controller service + runner.disableControllerService(livyControllerService); + } } diff --git a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractiveSSL.java b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractiveSSL.java index 1a379e4e151e..04808855cbbd 100644 --- a/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractiveSSL.java +++ b/nifi-nar-bundles/nifi-spark-bundle/nifi-livy-processors/src/test/java/org/apache/nifi/processors/livy/TestExecuteSparkInteractiveSSL.java @@ -17,6 +17,7 @@ package org.apache.nifi.processors.livy; import org.apache.nifi.controller.livy.LivySessionController; +import org.apache.nifi.controller.livy.utilities.LivyHelpers; import org.apache.nifi.ssl.StandardSSLContextService; import org.apache.nifi.util.TestRunners; import org.apache.nifi.web.util.TestServer; @@ -73,14 +74,11 @@ public void before() throws Exception { // Allow time for the controller service to fully initialize Thread.sleep(500); - LivySessionController livyControllerService = new LivySessionController(); + livyControllerService = new LivySessionController(); runner.addControllerService("livyCS", livyControllerService); - runner.setProperty(livyControllerService, LivySessionController.LIVY_HOST, url.substring(url.indexOf("://") + 3, url.lastIndexOf(":"))); - runner.setProperty(livyControllerService, LivySessionController.LIVY_PORT, url.substring(url.lastIndexOf(":") + 1)); - runner.setProperty(livyControllerService, LivySessionController.SSL_CONTEXT_SERVICE, "ssl-context"); - runner.enableControllerService(livyControllerService); - - runner.setProperty(ExecuteSparkInteractive.LIVY_CONTROLLER_SERVICE, "livyCS"); + runner.setProperty(livyControllerService, LivyHelpers.LIVY_HOST, url.substring(url.indexOf("://") + 3, url.lastIndexOf(":"))); + runner.setProperty(livyControllerService, LivyHelpers.LIVY_PORT, url.substring(url.lastIndexOf(":") + 1)); + runner.setProperty(livyControllerService, LivyHelpers.SSL_CONTEXT_SERVICE, "ssl-context"); server.clearHandlers(); }