service/src/main/java/org/apache/griffin/core/job/LivyTaskSubmitHelper.java - griffin - Git at Google

 /*
 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.griffin.core.job;

 import static org.apache.griffin.core.config.PropertiesConfig.livyConfMap;
 import static org.apache.griffin.core.job.entity.LivySessionStates.State.NOT_FOUND;
 import static org.apache.griffin.core.util.JsonUtil.toEntity;
 import static org.apache.griffin.core.util.JsonUtil.toJsonWithFormat;

 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;

 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.PostConstruct;

 import org.apache.commons.collections.map.HashedMap;
 import org.quartz.JobDetail;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.core.env.Environment;
 import org.springframework.http.HttpEntity;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.security.kerberos.client.KerberosRestTemplate;
 import org.springframework.stereotype.Component;
 import org.springframework.web.client.HttpClientErrorException;
 import org.springframework.web.client.RestTemplate;

 @Component
 public class LivyTaskSubmitHelper {
     private static final Logger LOGGER = LoggerFactory.getLogger(LivyTaskSubmitHelper.class);
     private static final String REQUEST_BY_HEADER = "X-Requested-By";
     public static final int DEFAULT_QUEUE_SIZE = 20000;
     private static final int SLEEP_TIME = 300;

     private SparkSubmitJob sparkSubmitJob;
     private ConcurrentMap<Long, Integer> taskAppIdMap = new ConcurrentHashMap<>();
     // Current number of tasks
     private AtomicInteger curConcurrentTaskNum = new AtomicInteger(0);
     private String workerNamePre;
     private RestTemplate restTemplate = new RestTemplate();
     // queue for pub or sub
     private BlockingQueue<JobDetail> queue;
     private String uri;

     @Value("${livy.task.max.concurrent.count:20}")
     private int maxConcurrentTaskCount;
     @Value("${livy.task.submit.interval.second:3}")
     private int batchIntervalSecond;

     @Autowired
     private Environment env;

     /**
      * Initialize related parameters and open consumer threads.
      */
     @PostConstruct
     public void init() {
         startWorker();
         uri = env.getProperty("livy.uri");
         LOGGER.info("Livy uri : {}", uri);
     }

     public LivyTaskSubmitHelper() {
         this.workerNamePre = "livy-task-submit-worker";
     }

     /**
      * Initialize blocking queues and start consumer threads.
      */
     public void startWorker() {
         queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_SIZE);
         ExecutorService executorService = Executors.newSingleThreadExecutor();
         TaskInner taskInner = new TaskInner(executorService);
         executorService.execute(taskInner);
     }

     /**
      * Put job detail into the queue.
      *
      * @param jd job detail.
      */
     public void addTaskToWaitingQueue(JobDetail jd) throws IOException {
         if (jd == null) {
             LOGGER.warn("task is blank, workerNamePre: {}", workerNamePre);
             return;
         }

         if (queue.remainingCapacity() <= 0) {
             LOGGER.warn("task is discard, workerNamePre: {}, task: {}", workerNamePre, jd);
             sparkSubmitJob.saveJobInstance(null, NOT_FOUND);
             return;
         }

         queue.add(jd);
         LOGGER.info("add_task_to_waiting_queue_success, workerNamePre: {}, task: {}",
             workerNamePre, jd);
     }

     /**
      * Consumer thread.
      */
     class TaskInner implements Runnable {
         private ExecutorService es;

         public TaskInner(ExecutorService es) {
             this.es = es;
         }

         public void run() {
             long insertTime = System.currentTimeMillis();
             while (true) {
                 try {
                     if (curConcurrentTaskNum.get() < maxConcurrentTaskCount
                         && (System.currentTimeMillis() - insertTime) >= batchIntervalSecond * 1000) {
                         JobDetail jd = queue.take();
                         sparkSubmitJob.saveJobInstance(jd);
                         insertTime = System.currentTimeMillis();
                     } else {
                         Thread.sleep(SLEEP_TIME);
                     }
                 } catch (Exception e) {
                     LOGGER.error("Async_worker_doTask_failed, {}", e.getMessage(), e);
                     es.execute(this);
                 }
             }
         }
     }

     /**
      * Add the batch id returned by Livy.
      *
      * @param scheduleId livy batch id.
      */
     public void increaseCurTaskNum(Long scheduleId) {
         curConcurrentTaskNum.incrementAndGet();
         if (scheduleId != null) {
             taskAppIdMap.put(scheduleId, 1);
         }
     }

     /**
      * Remove tasks after job status updates.
      *
      * @param scheduleId livy batch id.
      */
     public void decreaseCurTaskNum(Long scheduleId) {
         if (scheduleId != null && taskAppIdMap.containsKey(scheduleId)) {
             curConcurrentTaskNum.decrementAndGet();
             taskAppIdMap.remove(scheduleId);
         }
     }

     protected Map<String, Object> retryLivyGetAppId(String result, int appIdRetryCount)
         throws IOException {

         int retryCount = appIdRetryCount;
         TypeReference<HashMap<String, Object>> type =
             new TypeReference<HashMap<String, Object>>() {
             };
         Map<String, Object> resultMap = toEntity(result, type);

         if (retryCount <= 0) {
             return null;
         }

         if (resultMap.get("appId") != null) {
             return resultMap;
         }

         Object livyBatchesId = resultMap.get("id");
         if (livyBatchesId == null) {
             return resultMap;
         }

         while (retryCount-- > 0) {
             try {
                 Thread.sleep(SLEEP_TIME);
             } catch (InterruptedException e) {
                 LOGGER.error(e.getMessage(), e);
             }
             resultMap = getResultByLivyId(livyBatchesId, type);
             LOGGER.info("retry get livy resultMap: {}, batches id : {}", resultMap, livyBatchesId);

             if (resultMap.get("appId") != null) {
                 break;
             }
         }

         return resultMap;
     }

     private Map<String, Object> getResultByLivyId(Object livyBatchesId, TypeReference<HashMap<String, Object>> type)
         throws IOException {
         Map<String, Object> resultMap = new HashedMap();
         String livyUri = uri + "/" + livyBatchesId;
         String result = getFromLivy(livyUri);
         LOGGER.info(result);
         return result == null ? resultMap : toEntity(result, type);
     }

     public String postToLivy(String uri) {
         LOGGER.info("Post To Livy URI is: " + uri);
         String needKerberos = env.getProperty("livy.need.kerberos");
         LOGGER.info("Need Kerberos:" + needKerberos);

         HttpHeaders headers = new HttpHeaders();
         headers.setContentType(MediaType.APPLICATION_JSON);
         headers.set(REQUEST_BY_HEADER, "admin");

         if (needKerberos == null || needKerberos.isEmpty()) {
             LOGGER.error("The property \"livy.need.kerberos\" is empty");
             return null;
         }

         if (needKerberos.equalsIgnoreCase("false")) {
             LOGGER.info("The livy server doesn't need Kerberos Authentication");
             String result = null;
             try {
                 HttpEntity<String> springEntity = new HttpEntity<>(toJsonWithFormat(livyConfMap), headers);
                 result = restTemplate.postForObject(uri, springEntity, String.class);
                 LOGGER.info(result);
             } catch (HttpClientErrorException e) {
                 LOGGER.error("Post to livy ERROR. \n  response status : " + e.getMessage()
                     + "\n  response header : " + e.getResponseHeaders()
                     + "\n  response body : " + e.getResponseBodyAsString());
             } catch (JsonProcessingException e) {
                 LOGGER.error("Json Parsing failed, {}", e.getMessage(), e);
             } catch (Exception e) {
                 LOGGER.error("Post to livy ERROR. \n {}", e);
             }
             return result;
         } else {
             LOGGER.info("The livy server needs Kerberos Authentication");
             String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
             String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
             LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

             KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
             HttpEntity<String> springEntity = null;
             try {
                 springEntity = new HttpEntity<>(toJsonWithFormat(livyConfMap), headers);
             } catch (HttpClientErrorException e) {
                 LOGGER.error("Post to livy ERROR. \n  response status : " + e.getMessage()
                     + "\n  response header : " + e.getResponseHeaders()
                     + "\n  response body : " + e.getResponseBodyAsString());
             } catch (JsonProcessingException e) {
                 LOGGER.error("Json Parsing failed, {}", e.getMessage(), e);
             } catch (Exception e) {
                 LOGGER.error("Post to livy ERROR. {}", e.getMessage(), e);
             }
             String result = restTemplate.postForObject(uri, springEntity, String.class);
             LOGGER.info(result);
             return result;
         }
     }

     public String getFromLivy(String uri) {
         LOGGER.info("Get From Livy URI is: " + uri);
         String needKerberos = env.getProperty("livy.need.kerberos");
         LOGGER.info("Need Kerberos:" + needKerberos);

         if (needKerberos == null || needKerberos.isEmpty()) {
             LOGGER.error("The property \"livy.need.kerberos\" is empty");
             return null;
         }

         if (needKerberos.equalsIgnoreCase("false")) {
             LOGGER.info("The livy server doesn't need Kerberos Authentication");
             return restTemplate.getForObject(uri, String.class);
         } else {
             LOGGER.info("The livy server needs Kerberos Authentication");
             String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
             String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
             LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

             KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
             String result = restTemplate.getForObject(uri, String.class);
             LOGGER.info(result);
             return result;
         }
     }

     public void deleteByLivy(String uri) {
         LOGGER.info("Delete by Livy URI is: " + uri);
         String needKerberos = env.getProperty("livy.need.kerberos");
         LOGGER.info("Need Kerberos:" + needKerberos);

         if (needKerberos == null || needKerberos.isEmpty()) {
             LOGGER.error("The property \"livy.need.kerberos\" is empty");
             return;
         }

         if (needKerberos.equalsIgnoreCase("false")) {
             LOGGER.info("The livy server doesn't need Kerberos Authentication");
             new RestTemplate().delete(uri);
         } else {
             LOGGER.info("The livy server needs Kerberos Authentication");
             String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
             String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
             LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

             KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
             restTemplate.delete(uri);
         }
     }
 }
	/*
	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.griffin.core.job;

	import static org.apache.griffin.core.config.PropertiesConfig.livyConfMap;
	import static org.apache.griffin.core.job.entity.LivySessionStates.State.NOT_FOUND;
	import static org.apache.griffin.core.util.JsonUtil.toEntity;
	import static org.apache.griffin.core.util.JsonUtil.toJsonWithFormat;

	import com.fasterxml.jackson.core.JsonProcessingException;
	import com.fasterxml.jackson.core.type.TypeReference;

	import java.io.IOException;
	import java.util.HashMap;
	import java.util.Map;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.ConcurrentHashMap;
	import java.util.concurrent.ConcurrentMap;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.atomic.AtomicInteger;
	import javax.annotation.PostConstruct;

	import org.apache.commons.collections.map.HashedMap;
	import org.quartz.JobDetail;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;
	import org.springframework.beans.factory.annotation.Autowired;
	import org.springframework.beans.factory.annotation.Value;
	import org.springframework.core.env.Environment;
	import org.springframework.http.HttpEntity;
	import org.springframework.http.HttpHeaders;
	import org.springframework.http.MediaType;
	import org.springframework.security.kerberos.client.KerberosRestTemplate;
	import org.springframework.stereotype.Component;
	import org.springframework.web.client.HttpClientErrorException;
	import org.springframework.web.client.RestTemplate;

	@Component
	public class LivyTaskSubmitHelper {
	private static final Logger LOGGER = LoggerFactory.getLogger(LivyTaskSubmitHelper.class);
	private static final String REQUEST_BY_HEADER = "X-Requested-By";
	public static final int DEFAULT_QUEUE_SIZE = 20000;
	private static final int SLEEP_TIME = 300;

	private SparkSubmitJob sparkSubmitJob;
	private ConcurrentMap<Long, Integer> taskAppIdMap = new ConcurrentHashMap<>();
	// Current number of tasks
	private AtomicInteger curConcurrentTaskNum = new AtomicInteger(0);
	private String workerNamePre;
	private RestTemplate restTemplate = new RestTemplate();
	// queue for pub or sub
	private BlockingQueue<JobDetail> queue;
	private String uri;

	@Value("${livy.task.max.concurrent.count:20}")
	private int maxConcurrentTaskCount;
	@Value("${livy.task.submit.interval.second:3}")
	private int batchIntervalSecond;

	@Autowired
	private Environment env;

	/**
	* Initialize related parameters and open consumer threads.
	*/
	@PostConstruct
	public void init() {
	startWorker();
	uri = env.getProperty("livy.uri");
	LOGGER.info("Livy uri : {}", uri);
	}

	public LivyTaskSubmitHelper() {
	this.workerNamePre = "livy-task-submit-worker";
	}

	/**
	* Initialize blocking queues and start consumer threads.
	*/
	public void startWorker() {
	queue = new LinkedBlockingQueue<>(DEFAULT_QUEUE_SIZE);
	ExecutorService executorService = Executors.newSingleThreadExecutor();
	TaskInner taskInner = new TaskInner(executorService);
	executorService.execute(taskInner);
	}

	/**
	* Put job detail into the queue.
	*
	* @param jd job detail.
	*/
	public void addTaskToWaitingQueue(JobDetail jd) throws IOException {
	if (jd == null) {
	LOGGER.warn("task is blank, workerNamePre: {}", workerNamePre);
	return;
	}

	if (queue.remainingCapacity() <= 0) {
	LOGGER.warn("task is discard, workerNamePre: {}, task: {}", workerNamePre, jd);
	sparkSubmitJob.saveJobInstance(null, NOT_FOUND);
	return;
	}

	queue.add(jd);
	LOGGER.info("add_task_to_waiting_queue_success, workerNamePre: {}, task: {}",
	workerNamePre, jd);
	}

	/**
	* Consumer thread.
	*/
	class TaskInner implements Runnable {
	private ExecutorService es;

	public TaskInner(ExecutorService es) {
	this.es = es;
	}

	public void run() {
	long insertTime = System.currentTimeMillis();
	while (true) {
	try {
	if (curConcurrentTaskNum.get() < maxConcurrentTaskCount
	&& (System.currentTimeMillis() - insertTime) >= batchIntervalSecond * 1000) {
	JobDetail jd = queue.take();
	sparkSubmitJob.saveJobInstance(jd);
	insertTime = System.currentTimeMillis();
	} else {
	Thread.sleep(SLEEP_TIME);
	}
	} catch (Exception e) {
	LOGGER.error("Async_worker_doTask_failed, {}", e.getMessage(), e);
	es.execute(this);
	}
	}
	}
	}

	/**
	* Add the batch id returned by Livy.
	*
	* @param scheduleId livy batch id.
	*/
	public void increaseCurTaskNum(Long scheduleId) {
	curConcurrentTaskNum.incrementAndGet();
	if (scheduleId != null) {
	taskAppIdMap.put(scheduleId, 1);
	}
	}

	/**
	* Remove tasks after job status updates.
	*
	* @param scheduleId livy batch id.
	*/
	public void decreaseCurTaskNum(Long scheduleId) {
	if (scheduleId != null && taskAppIdMap.containsKey(scheduleId)) {
	curConcurrentTaskNum.decrementAndGet();
	taskAppIdMap.remove(scheduleId);
	}
	}

	protected Map<String, Object> retryLivyGetAppId(String result, int appIdRetryCount)
	throws IOException {

	int retryCount = appIdRetryCount;
	TypeReference<HashMap<String, Object>> type =
	new TypeReference<HashMap<String, Object>>() {
	};
	Map<String, Object> resultMap = toEntity(result, type);

	if (retryCount <= 0) {
	return null;
	}

	if (resultMap.get("appId") != null) {
	return resultMap;
	}

	Object livyBatchesId = resultMap.get("id");
	if (livyBatchesId == null) {
	return resultMap;
	}

	while (retryCount-- > 0) {
	try {
	Thread.sleep(SLEEP_TIME);
	} catch (InterruptedException e) {
	LOGGER.error(e.getMessage(), e);
	}
	resultMap = getResultByLivyId(livyBatchesId, type);
	LOGGER.info("retry get livy resultMap: {}, batches id : {}", resultMap, livyBatchesId);

	if (resultMap.get("appId") != null) {
	break;
	}
	}

	return resultMap;
	}

	private Map<String, Object> getResultByLivyId(Object livyBatchesId, TypeReference<HashMap<String, Object>> type)
	throws IOException {
	Map<String, Object> resultMap = new HashedMap();
	String livyUri = uri + "/" + livyBatchesId;
	String result = getFromLivy(livyUri);
	LOGGER.info(result);
	return result == null ? resultMap : toEntity(result, type);
	}

	public String postToLivy(String uri) {
	LOGGER.info("Post To Livy URI is: " + uri);
	String needKerberos = env.getProperty("livy.need.kerberos");
	LOGGER.info("Need Kerberos:" + needKerberos);

	HttpHeaders headers = new HttpHeaders();
	headers.setContentType(MediaType.APPLICATION_JSON);
	headers.set(REQUEST_BY_HEADER, "admin");

	if (needKerberos == null \|\| needKerberos.isEmpty()) {
	LOGGER.error("The property \"livy.need.kerberos\" is empty");
	return null;
	}

	if (needKerberos.equalsIgnoreCase("false")) {
	LOGGER.info("The livy server doesn't need Kerberos Authentication");
	String result = null;
	try {
	HttpEntity<String> springEntity = new HttpEntity<>(toJsonWithFormat(livyConfMap), headers);
	result = restTemplate.postForObject(uri, springEntity, String.class);
	LOGGER.info(result);
	} catch (HttpClientErrorException e) {
	LOGGER.error("Post to livy ERROR. \n response status : " + e.getMessage()
	+ "\n response header : " + e.getResponseHeaders()
	+ "\n response body : " + e.getResponseBodyAsString());
	} catch (JsonProcessingException e) {
	LOGGER.error("Json Parsing failed, {}", e.getMessage(), e);
	} catch (Exception e) {
	LOGGER.error("Post to livy ERROR. \n {}", e);
	}
	return result;
	} else {
	LOGGER.info("The livy server needs Kerberos Authentication");
	String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
	String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
	LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

	KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
	HttpEntity<String> springEntity = null;
	try {
	springEntity = new HttpEntity<>(toJsonWithFormat(livyConfMap), headers);
	} catch (HttpClientErrorException e) {
	LOGGER.error("Post to livy ERROR. \n response status : " + e.getMessage()
	+ "\n response header : " + e.getResponseHeaders()
	+ "\n response body : " + e.getResponseBodyAsString());
	} catch (JsonProcessingException e) {
	LOGGER.error("Json Parsing failed, {}", e.getMessage(), e);
	} catch (Exception e) {
	LOGGER.error("Post to livy ERROR. {}", e.getMessage(), e);
	}
	String result = restTemplate.postForObject(uri, springEntity, String.class);
	LOGGER.info(result);
	return result;
	}
	}

	public String getFromLivy(String uri) {
	LOGGER.info("Get From Livy URI is: " + uri);
	String needKerberos = env.getProperty("livy.need.kerberos");
	LOGGER.info("Need Kerberos:" + needKerberos);

	if (needKerberos == null \|\| needKerberos.isEmpty()) {
	LOGGER.error("The property \"livy.need.kerberos\" is empty");
	return null;
	}

	if (needKerberos.equalsIgnoreCase("false")) {
	LOGGER.info("The livy server doesn't need Kerberos Authentication");
	return restTemplate.getForObject(uri, String.class);
	} else {
	LOGGER.info("The livy server needs Kerberos Authentication");
	String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
	String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
	LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

	KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
	String result = restTemplate.getForObject(uri, String.class);
	LOGGER.info(result);
	return result;
	}
	}

	public void deleteByLivy(String uri) {
	LOGGER.info("Delete by Livy URI is: " + uri);
	String needKerberos = env.getProperty("livy.need.kerberos");
	LOGGER.info("Need Kerberos:" + needKerberos);

	if (needKerberos == null \|\| needKerberos.isEmpty()) {
	LOGGER.error("The property \"livy.need.kerberos\" is empty");
	return;
	}

	if (needKerberos.equalsIgnoreCase("false")) {
	LOGGER.info("The livy server doesn't need Kerberos Authentication");
	new RestTemplate().delete(uri);
	} else {
	LOGGER.info("The livy server needs Kerberos Authentication");
	String userPrincipal = env.getProperty("livy.server.auth.kerberos.principal");
	String keyTabLocation = env.getProperty("livy.server.auth.kerberos.keytab");
	LOGGER.info("principal:{}, lcoation:{}", userPrincipal, keyTabLocation);

	KerberosRestTemplate restTemplate = new KerberosRestTemplate(keyTabLocation, userPrincipal);
	restTemplate.delete(uri);
	}
	}
	}