Spring boot + Websocket 初篇

dependencies { 
  compile('org.springframework.boot:spring-boot-starter-security')
  compile('org.springframework.boot:spring-boot-starter-thymeleaf')
  compile('org.springframework.boot:spring-boot-starter-websocket')
  compile('org.thymeleaf.extras:thymeleaf-extras-springsecurity4')
  compile('org.mybatis.spring.boot:mybatis-spring-boot-starter:1.3.0')
  runtime('mysql:mysql-connector-java')
  compile('org.springframework.boot:spring-boot-devtools')
}

package com.demo.entity;

import com.demo.mapper.CollectorMapper;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.List;
import java.util.TimerTask;

// 定时从数据库中捞数据
public class TimerDBTask extends TimerTask {
    private WebSocketSession socketSession;
    private CollectorMapper collectorMapper;
    private int total = 0;

    @Override
    public void run() {
        int latest = collectorMapper.countTotal();
        if(latest > total){
            List<CollectData> dataList = collectorMapper.getCollectedValues();
            StringBuilder initial = new StringBuilder();
            for(CollectData collectData: dataList){
                initial.append(collectData.toString());
                initial.append("|");
            }
            if(initial.length() > 1) {
                initial.deleteCharAt(initial.length() - 1);
            }
            total = latest;
            try {
                socketSession.sendMessage(new TextMessage(initial.toString()));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public TimerDBTask(WebSocketSession session, CollectorMapper collectorMapper){
        this.socketSession = session;
        this.collectorMapper = collectorMapper;
    }
}

package com.demo.config;

import com.demo.entity.MessageHandler;
import com.demo.mapper.CollectorMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebsocketConfig implements WebSocketConfigurer{
    @Autowired
    private CollectorMapper collectorMapper;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(messageHandler(), "/websocket-endpoint");
    }

    @Bean
    public WebSocketHandler messageHandler(){
        return new MessageHandler(this.collectorMapper);
    }
}

package com.demo.entity;

import com.demo.mapper.CollectorMapper;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Timer;

public class MessageHandler extends TextWebSocketHandler {
    private Timer timer;
    private CollectorMapper collectorMapper;

    public MessageHandler(){

    }

    public MessageHandler(CollectorMapper collectorMapper){
        this.collectorMapper = collectorMapper;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // The websocket has been opened
        // might save this session object so that can send messages to it outside of this method
        String loginUser = session.getPrincipal().getName();
        timer = new Timer(true);
        long delay = 0l;
        TimerDBTask dbTask = new TimerDBTask(session, collectorMapper);
        timer.schedule(dbTask, delay, 60000);
//        session.sendMessage(new TextMessage("You're now connected to the server.This is the first message."));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        super.handleMessage(session, message);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // a message has been received
        System.out.println("Message received:" + message.getPayload());
        if(!session.isOpen()){
            timer.cancel();
            return;
        }
        super.handleTextMessage(session, message);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        super.afterConnectionClosed(session, status);
        // The websocket has been closed
    }
}