spark连接kafka工具类

import kafka.common.TopicAndPartition
import kafka.message.MessageAndMetadata
import kafka.serializer.Decoder
import org.apache.spark.SparkException
import org.apache.spark.rdd.RDD
import org.apache.spark.streaming.StreamingContext
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka.KafkaCluster.LeaderOffset
import org.apache.spark.streaming.kafka.{HasOffsetRanges, KafkaCluster, KafkaUtils, OffsetRange}
import org.slf4j
import org.slf4j.LoggerFactory
import redis.clients.jedis.{JedisCluster, Response}

import scala.collection.JavaConversions._
import scala.collection.mutable.ArrayBuffer
import scala.reflect.ClassTag
import scala.util.control.Breaks.breakable

/**
  * Create by shengjk1 on 2018/1/3
  */

class KafkaManager(var kafkaParams: Map[String, String]) extends Serializable {
	@transient
	val logger: slf4j.Logger = LoggerFactory.getLogger(this.getClass)
	
	val offsetStoreInRedisDayCounts=3
	
	kafkaParams+=("refresh.leader.backoff.ms"->"3000")
	val kc = new KafkaCluster(kafkaParams)
	
	/**
	  *
	  * @param ssc
	  * @param kafkaParams
	  * @param topicPartition Set([test01,0], [test01,1], [test01,2]))
	  * @tparam K
	  * @tparam V
	  * @tparam KD
	  * @tparam VD
	  * @return
	  */
	def createDirectStreamByAssignPartition[K: ClassTag, V: ClassTag, KD <: Decoder[K] : ClassTag, VD <: Decoder[V] : ClassTag](
					ssc: StreamingContext, kafkaParams: Map[String, String], topicPartition: Set[TopicAndPartition]): InputDStream[(K, V)] = {
	
		val groupId = kafkaParams.get("group.id").get
		
		setOrUpdateOffsetsByAssignPartition(topicPartition, groupId)
		//从zookeeper上读取offset开始消费message
		val messages = {
			//TODO getfromoffset
			val consumerOffsetsE = kc.getConsumerOffsets(groupId, topicPartition)//Right(Map([order5test1,2] -> 77262))
			if (consumerOffsetsE.isLeft)
				throw new SparkException(s"get kafka consumer offsets failed: ${consumerOffsetsE.left.get}")
			val consumerOffsets = consumerOffsetsE.right.get //([order5test1,2],77262)
			val offsets: Map[TopicAndPartition, Long] = consumerOffsets //([order5test1,2],77262)
			KafkaUtils.createDirectStream[K, V, KD, VD, (K, V)](
				ssc, kafkaParams, consumerOffsets, (mmd: MessageAndMetadata[K, V]) => (mmd.key, mmd.message))
		}
		messages
	}
	
	
	
	/**
	  * 从指定offset消费kafka
	  *
	  * @param ssc
	  * @param kafkaParams
	  * @param list
	  * @tparam K
	  * @tparam V
	  * @tparam KD
	  * @tparam VD
	  * @return
	  */
	def createDirectStreamByAssignPartitionAndOffset[K: ClassTag, V: ClassTag, KD <: Decoder[K] : ClassTag, VD <: Decoder[V] : ClassTag](
								ssc: StreamingContext, kafkaParams: Map[String, String], list: List[(String, Int, Long)]): InputDStream[(K, V)] = {
		
		val fromOffsets: Map[TopicAndPartition, Long] = setFromOffsets(list)
		
		val groupId = kafkaParams.get("group.id").get
		// 在zookeeper上读取offsets前先根据实际情况更新offsets
		val topics = ArrayBuffer[String]()
		val topicPartition = fromOffsets.keys.toSet	setOrUpdateOffsetsByAssignPartition(topicPartition, groupId)
		//从kafka上指定位置的offset开始消费message
		val messages = {
			KafkaUtils.createDirectStream[K, V, KD, VD, (K, V)](
				ssc, kafkaParams, fromOffsets, (mmd: MessageAndMetadata[K, V]) => (mmd.key, mmd.message))
		}
		messages
	}
	
	
	/**
	  * 创建数据流前，根据实际消费情况更新消费offsets
	  *
	  * @param topicPartition
	  * @param groupId
	  */
	private def setOrUpdateOffsetsByAssignPartition(topicPartition: Set[TopicAndPartition], groupId: String): Unit = {
		var hasConsumed = true
		val partitions = topicPartition //[order5test1,2]
		val consumerOffsetsE = kc.getConsumerOffsets(groupId, partitions)//Right(Map([order5test1,2] -> 77262))
		if (consumerOffsetsE.isLeft) hasConsumed = false
		if (hasConsumed) {
			// 消费过
			/**
			  * 如果streaming程序执行的时候出现kafka.common.OffsetOutOfRangeException，
			  * 说明zk上保存的offsets已经过时了，即kafka的定时清理策略已经将包含该offsets的文件删除。
			  * 针对这种情况，只要判断一下zk上的consumerOffsets和earliestLeaderOffsets的大小，
			  * 如果consumerOffsets比earliestLeaderOffsets还小的话，说明consumerOffsets已过时,
			  * 这时把consumerOffsets更新为earliestLeaderOffsets
			  */
			
			// 可能只是存在部分分区consumerOffsets过时，所以只更新过时分区的consumerOffsets为earliestLeaderOffsets
			val earliestLeaderOffsetsE = kc.getEarliestLeaderOffsets(partitions)//Right(Map([order5test1,2] -> LeaderOffset(local,9092,16556)))
			if (earliestLeaderOffsetsE.isLeft)
				throw new SparkException(s"get earliest leader offsets failed: ${earliestLeaderOffsetsE.left.get}")
			val earliestLeaderOffsets = earliestLeaderOffsetsE.right.get //([order5test1,2],LeaderOffset(local,9092,16556))
			val consumerOffsets = consumerOffsetsE.right.get//([order5test1,2],77262)
			
			// 可能只是存在部分分区consumerOffsets过时，所以只更新过时分区的consumerOffsets为earliestLeaderOffsets
			var offsets: Map[TopicAndPartition, Long] = Map()
			consumerOffsets.foreach({ case (tp, n) => //n 77262
				val earliestLeaderOffset = earliestLeaderOffsets(tp).offset//16556
				if (n < earliestLeaderOffset) {
					logger.info("consumer group:" + groupId + ",topic:" + tp.topic + ",partition:" + tp.partition +
							" offsets已经过时，更新为" + earliestLeaderOffset)
					offsets += (tp -> earliestLeaderOffset)
				}
			})
			if (!offsets.isEmpty) {
				kc.setConsumerOffsets(groupId, offsets)
			}
			
			
			//右越界 当consumerOffset>latesLeaderOffsets 程序员自己处理
			//				val latesLeaderOffsetsE = kc.getLatestLeaderOffsets(partitions)
			
		} else {
			// 没有消费过
			val reset = kafkaParams.get("auto.offset.reset").map(_.toLowerCase)
			var leaderOffsets: Map[TopicAndPartition, LeaderOffset] = null
			if (reset == Some("smallest")) {
				val leaderOffsetsE = kc.getEarliestLeaderOffsets(partitions)
				if (leaderOffsetsE.isLeft)
					throw new SparkException(s"get earliest leader offsets failed: ${leaderOffsetsE.left.get}")
				leaderOffsets = leaderOffsetsE.right.get
			} else {
				val leaderOffsetsE = kc.getLatestLeaderOffsets(partitions)
				if (leaderOffsetsE.isLeft)
					throw new SparkException(s"get latest leader offsets failed: ${leaderOffsetsE.left.get}")
				leaderOffsets = leaderOffsetsE.right.get
			}
			val offsets = leaderOffsets.map {
				case (tp, offset) => (tp, offset.offset)
			}
			kc.setConsumerOffsets(groupId, offsets)
		}
	}
	
	
	
	/**
	  * 构建Map
	  *
	  * @param list string:topicName Int:分区号 Long:指定偏移量
	  * @return
	  */
	def setFromOffsets(list: List[(String, Int, Long)]): Map[TopicAndPartition, Long] = {
		var fromOffsets: Map[TopicAndPartition, Long] = Map()
		for (offset <- list) {
			val tp = TopicAndPartition(offset._1, offset._2) //topic和分区数
			fromOffsets += (tp -> offset._3) // offset位置
		}
		fromOffsets
	}
	
	
	/**
	  * 创建数据流前，根据实际消费情况更新消费offsets
	  * @param topics
	  * @param groupId
	  */
	private def setOrUpdateOffsets(topics: Set[String], groupId: String): Unit = {
		topics.foreach(topic => {
			var hasConsumed = true
			val partitionsE = kc.getPartitions(Set(topic))
			if (partitionsE.isLeft)
				throw new SparkException(s"get kafka partition failed: ${partitionsE.left.get}")
			val partitions = partitionsE.right.get
			val consumerOffsetsE = kc.getConsumerOffsets(groupId, partitions)
			if (consumerOffsetsE.isLeft) hasConsumed = false
			if (hasConsumed) {
				// 消费过
				/**
				  *
				  * 如果消息体太大了，超过 fetch.message.max.bytes=1m的默认配置，
				  * 那么Spark Streaming会直接抛出OffsetOutOfRangeException异常，然后停止服务。
				  * 解决方案：Kafka consumer中设置fetch.message.max.bytes为大一点的内存
				  *
				  * 如果streaming程序执行的时候出现kafka.common.OffsetOutOfRangeException，
				  * 说明zk上保存的offsets已经过时了，即kafka的定时清理策略已经将包含该offsets的文件删除。
				  * 针对这种情况，只要判断一下zk上的consumerOffsets和earliestLeaderOffsets的大小，
				  * 如果consumerOffsets比earliestLeaderOffsets还小的话，说明consumerOffsets已过时,
				  * 这时把consumerOffsets更新为earliestLeaderOffsets
				  */
					
					//左越界 当consumerOffset<earliesLeaderOffsets时，置consumerOffset为earliesLeaderOffsets
				val earliestLeaderOffsetsE = kc.getEarliestLeaderOffsets(partitions)
				if (earliestLeaderOffsetsE.isLeft)
					throw new SparkException(s"get earliest leader offsets failed: ${earliestLeaderOffsetsE.left.get}")
				val earliestLeaderOffsets = earliestLeaderOffsetsE.right.get
				val consumerOffsets = consumerOffsetsE.right.get
				
				// 可能只是存在部分分区consumerOffsets过时，所以只更新过时分区的consumerOffsets为earliestLeaderOffsets
				var offsets: Map[TopicAndPartition, Long] = Map()
				consumerOffsets.foreach({ case (tp, n) =>
					val earliestLeaderOffset = earliestLeaderOffsets(tp).offset
					if (n < earliestLeaderOffset) {
						logger.info("consumer group:" + groupId + ",topic:" + tp.topic + ",partition:" + tp.partition +
								" offsets已经过时，更新为" + earliestLeaderOffset)
						offsets += (tp -> earliestLeaderOffset)
					}
				})
				if (!offsets.isEmpty) {
					kc.setConsumerOffsets(groupId, offsets)
				}
				
				
				//右越界 当consumerOffset>latesLeaderOffsets 程序员自己处理
//				val latesLeaderOffsetsE = kc.getLatestLeaderOffsets(partitions)
				
				
			} else {
				// 没有消费过
				val reset = kafkaParams.get("auto.offset.reset").map(_.toLowerCase)
				var leaderOffsets: Map[TopicAndPartition, LeaderOffset] = null
				if (reset == Some("smallest")) {
					val leaderOffsetsE = kc.getEarliestLeaderOffsets(partitions)
					if (leaderOffsetsE.isLeft)
						throw new SparkException(s"get earliest leader offsets failed: ${leaderOffsetsE.left.get}")
					leaderOffsets = leaderOffsetsE.right.get
				} else {
					val leaderOffsetsE = kc.getLatestLeaderOffsets(partitions)
					if (leaderOffsetsE.isLeft)
						throw new SparkException(s"get latest leader offsets failed: ${leaderOffsetsE.left.get}")
					leaderOffsets = leaderOffsetsE.right.get
				}
				val offsets = leaderOffsets.map {
					case (tp, offset) => (tp, offset.offset)
				}
				kc.setConsumerOffsets(groupId, offsets)
			}
		})
	}
	
	
	/**
	  * 更新zookeeper上的消费offsets
	  *
	  * @param rdd
	  */
	def updateZKOffsets(rdd: RDD[(String, String)]): Unit = {
		val groupId = kafkaParams.get("group.id").get
		val offsetsList = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
		
		for (offsets <- offsetsList) {
			val topicAndPartition = TopicAndPartition(offsets.topic, offsets.partition)
			val o = kc.setConsumerOffsets(groupId, Map((topicAndPartition, offsets.untilOffset)))
			if (o.isLeft) {
				throw new SparkException(s"Error updating the offset to Kafka cluster: ${o.left.get}")
			}
		}
	}
	
	
	/**
	  * 更新offsets并缓存redis
	  *
	  * @param rdd
	  */
	def updateZKOffsetsAndStoreRedis(rdd: RDD[(String, String)],conn:JedisCluster,key:String): Unit = {
		val groupId = kafkaParams.get("group.id").get
		val offsetsList = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
		
		val pipeline = JedisClusterPipeline.pipelined(conn)
		pipeline.refreshCluster()
		
		for (offsets <- offsetsList) {
			val topicAndPartition = TopicAndPartition(offsets.topic, offsets.partition)
			val o = kc.setConsumerOffsets(groupId, Map((topicAndPartition, offsets.untilOffset)))
			if (o.isLeft) {
				throw new SparkException(s"Error updating the offset to Kafka cluster: ${o.left.get}")
			}
			//			println(DateUtils.getNowEndHour+"-"+groupId+"-"+offsets.topic+"-"+offsets.partition+"#"+offsets.untilOffset)
			val listStr:java.util.List[String]=conn.lrange(key,0,-1)
			
			if(listStr.size()>0 &&listStr(0).contains(DateUtils.getNowDay("yyyyMMdd HH"))){
				pipeline.lpop(key)
				pipeline.lpush(key,listStr(0).split("#")(0)+"#"+offsets.untilOffset)
			}else{
				pipeline.lpush(key,DateUtils.getNowDay("yyyyMMdd HH")+"-"+groupId+"-"+offsets.topic+"-"+offsets.partition+"#"+offsets.untilOffset)
			}
			pipeline.expire(key,Constant.REDIS_OFFSET_EXPIRE)
			
			
			var offsetCount=listStr.size()
			while(offsetCount>offsetStoreInRedisDayCounts){
				pipeline.rpop(key)
				offsetCount=offsetCount-1
			}
		}
		pipeline.sync()
	}
	
	
	/**
	  * @param offsetsList
	  */
	def updateZKOffsets(offsetsList:Array[OffsetRange]): Unit = {
		val groupId = kafkaParams.get("group.id").get
		
		for (offsets <- offsetsList) {
			val topicAndPartition = TopicAndPartition(offsets.topic, offsets.partition)
			val o = kc.setConsumerOffsets(groupId, Map((topicAndPartition, offsets.untilOffset)))
			if (o.isLeft) {
				throw new SparkException(s"Error updating the offset to Kafka cluster: ${o.left.get}")
			}
		}
	}
}