Spark Streaming监听Socket解决自动断开连接

Spark Streaming与Socket

我们在使用Spark Streaming的实时计算功能的时候，有时候会与Socket进行通信，即通过Spark Streaming监听Socket，实时获取从Socket发送的数据进行处理。
但最近发现一个坑：Socket如果客户端与服务端之间有一段时间没有进行通信，即客户端与服务端之间相互没有进行数据传输，那么客户端与服务端的连接就会自动断开。
参考了一些资料，一般的做法就是在客户端增加一个“心跳机制”：客户端每隔一段时间往服务端发送一次数据，使得连接不会自动断开。在原生的Java或其他使用中，就很容易实现。就是起一个线程来不断运行这个心跳机制的方法。

自定义接收器

那么，在Spark Streaming中如何实现呢？
可能很多人像我之前一样，在Spark Streaming与Socket通信的时候，都是使用这个官方APIsocketTextStream。但是，你想要在增加“心跳机制”的话，就不能使用这个API，需要自定义接收器Spark Streaming Custom Receivers。
首先，你需要实现一个接收器类，然后继承org.apache.spark.streaming.receiver.Receiver

1. onStart()：开始接收数据的方法。在这个方法里面，你可以创建你的客户端连接，然后开启一个线程来接收数据，即执行receive方法
2. onStop()：停止接收数据的方法。在这个方法，一般不需要任何处理
3. receive()：接收数据的方法

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.ConnectException;
import java.net.Socket;
import java.nio.charset.StandardCharsets;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.VoidFunction;
import org.apache.spark.storage.StorageLevel;
import org.apache.spark.streaming.Duration;
import org.apache.spark.streaming.api.java.JavaReceiverInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.receiver.Receiver;

import com.google.common.io.Closeables;

public class JavaCustomReceiver extends Receiver<String> {

	private static final long serialVersionUID = 1L;

	public static void main(String[] args) throws Exception {

		// Create the context with a 1 second batch size
		SparkConf sparkConf = new SparkConf().setAppName("JavaCustomReceiver").setMaster("local[4]");
		JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, new Duration(3000));
		ssc.sparkContext().setLogLevel("ERROR");

		JavaReceiverInputDStream<String> lines = ssc.receiverStream(new JavaCustomReceiver("localhost", 9999));

		lines.foreachRDD(new VoidFunction<JavaRDD<String>>() {
			
			private static final long serialVersionUID = 1L;

			@Override
			public void call(JavaRDD<String> t) throws Exception {
				System.out.println("data -->" + t.collect());
			}
		});
		
		ssc.start();
		ssc.awaitTermination();
	}

	// ============= Receiver code that receives data over a socket ==============

	String host = null;
	int port = -1;
	private Socket socket = null;
	private boolean firstStart = true;

	public JavaCustomReceiver(String host_, int port_) {
		super(StorageLevel.MEMORY_AND_DISK_2());
		host = host_;
		port = port_;
	}
	
	/**
	 * 初设化或receiver异常的时候，会有重连机制，调用restart方法，然后执行该方法onStart
	 */
	@Override
	public void onStart() {
		// Start the thread that receives data over a connection
		try {
			socket = new Socket(host, port);
			socket.setKeepAlive(true);
		} catch (Exception e) {
			restart("Error reconnect to socket", e);
		}
		new Thread(this::receive).start();
		
		// 只需要在第一次开始的时候，创建一个线程就可以了。
		if (firstStart) {
			new Thread(this::heartBeat).start();
		}
	}
	
	@Override
	public void onStop() {
		// There is nothing much to do as the thread calling receive()
		// is designed to stop by itself isStopped() returns false
	}

	/** Create a socket connection and receive data until receiver is stopped */
	private void receive() {
		try {
			BufferedReader reader = null;
			try {
				// connect to the server
				reader = new BufferedReader(new InputStreamReader(socket.getInputStream(), StandardCharsets.UTF_8));
				// Until stopped or connection broken continue reading
				String userInput;
				// 当receiver正常工作时，监听线程会一直处在这个while循环运行
				while (!isStopped() && (userInput = reader.readLine()) != null) {
					// 当监听有数据时，就会进来，执行此处的代码块
					System.out.println("Received data '" + userInput + "'");
					store(userInput);
				}
			} finally {
				Closeables.close(reader, /* swallowIOException = */ true);
				Closeables.close(socket, /* swallowIOException = */ true);
			}
			// Restart in an attempt to connect again when server is active again
			restart("Trying to connect again");
		} catch (ConnectException ce) {
			// restart if could not connect to server
			restart("Could not connect", ce);
		} catch (Throwable t) {
			restart("Error receiving data", t);
		}
	}
	
	/**
	 * 心跳线程
	 */
	private void heartBeat() {
		firstStart = false;
		OutputStream out;
		while (true) {
			// 每隔10分钟往服务端，发送一次"心跳"，防止自动断开
			try {
				Thread.sleep(10 * 60 * 1000);
			} catch (InterruptedException e) {
				e.printStackTrace();
				break;
			}

			synchronized (socket) {
				if (socket != null && !socket.isClosed()) {
					try {
						out = socket.getOutputStream();
						out.write(0);
						out.flush();
						// 这个地方不能将out.close()，否则socket连接也会一同关闭，receiver接收数据就会报错
//						out.close(); 
						out = null;
						System.out.println("socket is closed: " + socket.isClosed());
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
				}
			}
		}
	}
}

运行机制

1. 首先，第一次的时候，会执行onStart方法，然后进行开始接收数据，进行receive方法；
2. 然后，会在receive方法中的while循环不断执行，没有数据的时候就卡在while的判断体，当有数据的时候，就会进入while里面的方法体，这里就必须有store方法来接收数据；
   
3. 当发生连接异常或其他异常的必须进行捕获，然后调用restart方法，restart方法稍后则会调用onStart，重新建立连接，监听并接收数据。

欢迎关注同名公众号：“我就算饿死也不做程序员”。
交个朋友，一起交流，一起学习，一起进步。