remove signal in reader

等待ping时间改为30秒
整理关闭逻辑
增加文档

serv.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"flag"
 	"net/http"
@@ -8,7 +9,6 @@ import (
 	"os/signal"
 	"strconv"
 	"strings"
-	"syscall"
 	"time"
 
 	"github.com/gorilla/handlers"
@@ -24,29 +24,62 @@ const (
 	writeWait = 10 * time.Second
 
 	// Time allowed to read the next pong message from the peer.
-	pongWait = 60 * time.Second
+	pongWait = 30 * time.Second
 
 	// Send pings to peer with this period. Must be less than pongWait.
 	pingPeriod = (pongWait * 9) / 10
 
 	// Maximum message size allowed from peer.
-	// maxMessageSize = 512
+	maxMessageSize = 512
 )
 
 var (
 	newline          = []byte{'\n'}
+	space            = []byte{' '}
 	addr             = flag.String("addr", ":9000", "http service address")
 	bootstrapServers = flag.String("bootstrap-servers", "localhost:9092", "kafka bootstrap servers")
 	// partition        = flag.Int("partition", 0, "kafka tipic partion")
 	wait     = flag.Duration("graceful-timeout", time.Second*15, "the duration for which the server gracefully wait for existing connections to finish - e.g. 15s or 1m")
 	upgrader = websocket.Upgrader{
-		ReadBufferSize:  1024,
-		WriteBufferSize: 1024,
+		ReadBufferSize:  2048,
+		WriteBufferSize: 2048,
 	}
 	hubMap = make(map[string]Hub)
 	router = mux.NewRouter()
 )
 
+// Hub maintains the kafka reader and set of active clients
+// broadcasts messages to the clients.
+type Hub struct {
+
+	// hun running indicator
+	running bool
+
+	// Kafka topic
+	topic string
+
+	// Kakfa reader partition
+	partition int
+
+	// Kafka start offset, -1 = latest, -2 = ealist
+	offset int64
+
+	// Kafka Reader
+	reader *kafka.Reader
+
+	// Registered clients.
+	clients map[*Client]bool
+
+	// Inbound messages from the clients.
+	broadcast chan []byte
+
+	// Register requests from the clients.
+	register chan *Client
+
+	// Unregister requests from clients.
+	unregister chan *Client
+}
+
 // Client is a middleman between the websocket connection and the hub.
 type Client struct {
 	hub *Hub
@@ -58,7 +91,9 @@ type Client struct {
 	send chan []byte
 }
 
-func create(topic string, offset int64, partition int) *kafka.Reader {
+func createKafkaReader(topic string, offset int64, partition int) *kafka.Reader {
+
+	// Use uuid for kafka group id
 	group := uuid.NewV4()
 
 	log.WithFields(log.Fields{
@@ -69,48 +104,54 @@ func create(topic string, offset int64, partition int) *kafka.Reader {
 		"offset":            offset}).Info("creating kafka client ... ")
 
 	reader := kafka.NewReader(kafka.ReaderConfig{
-		Brokers:   strings.Split(*bootstrapServers, ","),
-		Topic:     topic,
-		Partition: partition,
-		MinBytes:  10e1,
-		MaxBytes:  10e6, // 10MB
+		Brokers:   strings.Split(*bootstrapServers, ","), // kafka brokers, split by comma
+		Topic:     topic,                                 // topic to subscribe
+		Partition: partition,                             // partition from url
+		MinBytes:  10e2,                                  // 1k
+		MaxBytes:  10e6,                                  // 10MB
 	})
+
+	// offset from url
+	// -1 = latest
+	// -2 = ealiest
 	reader.SetOffset(offset)
 
 	return reader
 }
 
-// readPump pumps messages from the websocket connection to the hub.
-//
-// The application runs readPump in a per-connection goroutine. The application
-// ensures that there is at most one reader on a connection by executing all
-// reads from this goroutine.
-func (c *Client) readPump() {
+// startClient keep the client connection running
+// response the pong message
+
+func (c *Client) startClient() {
+
+	log.WithField("remote", c.conn.RemoteAddr()).Info("client started")
+
 	defer func() {
 		c.hub.unregister <- c
 		c.conn.Close()
 		log.WithField("client", c).Warn("exit")
 	}()
-	// c.conn.SetReadLimit(maxMessageSize)
+
+	// response to pong message
+	c.conn.SetReadLimit(maxMessageSize)
 	c.conn.SetReadDeadline(time.Now().Add(pongWait))
 	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
 
-	sigchan := make(chan os.Signal, 1)
-	signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)
-	run := true
-	for run == true {
-		select {
-		case sig := <-sigchan:
-			log.WithFields(log.Fields{"signal": sig}).Info(" terminating")
-			run = false
-			for topic, hub := range hubMap {
-				log.WithField("topic", topic).Warn("closing hub")
-				hub.running = false
+	for {
+		_, message, err := c.conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				log.WithError(err).Error("read message error")
 			}
-			log.Fatal("break ")
+			break
+		}
+		message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
+
+		// if client send ping, response pong
+		if string(message) == "ping" {
+			c.send <- []byte("pong")
 		}
 	}
-
 }
 
 // writePump pumps messages from the hub to the websocket connection.
@@ -119,6 +160,8 @@ func (c *Client) readPump() {
 // application ensures that there is at most one writer to a connection by
 // executing all writes from this goroutine.
 func (c *Client) writePump() {
+
+	// ping message
 	ticker := time.NewTicker(pingPeriod)
 	defer func() {
 		ticker.Stop()
@@ -126,11 +169,13 @@ func (c *Client) writePump() {
 		c.hub.unregister <- c
 		log.WithField("client", c).Warn("exit")
 	}()
+
 	for {
 		select {
 		case message, ok := <-c.send:
 			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
 			if !ok {
+				log.WithField("client", c).Warn("sending close message")
 				// The hub closed the channel.
 				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
 				return
@@ -138,6 +183,7 @@ func (c *Client) writePump() {
 
 			w, err := c.conn.NextWriter(websocket.TextMessage)
 			if err != nil {
+				log.WithError(err).Error("can't get next wirter")
 				return
 			}
 			w.Write(message)
@@ -150,11 +196,13 @@ func (c *Client) writePump() {
 			}
 
 			if err := w.Close(); err != nil {
+				log.WithError(err).Error("error in close writer")
 				return
 			}
 		case <-ticker.C:
 			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
 			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				log.WithError(err).Error("can't write ping message")
 				return
 			}
 		}
@@ -194,33 +242,7 @@ func websocketHandler(w http.ResponseWriter, r *http.Request) {
 	// Allow collection of memory referenced by the caller by doing all work in
 	// new goroutines.
 	go client.writePump()
-	go client.readPump()
-}
-
-// Hub maintains the set of active clients and broadcasts messages to the
-// clients.
-type Hub struct {
-	running bool
-
-	topic string
-
-	partition int
-
-	offset int64
-
-	reader *kafka.Reader
-
-	// Registered clients.
-	clients map[*Client]bool
-
-	// Inbound messages from the clients.
-	broadcast chan []byte
-
-	// Register requests from the clients.
-	register chan *Client
-
-	// Unregister requests from clients.
-	unregister chan *Client
+	go client.startClient()
 }
 
 func newHub(topic string, offset int64, partition int) *Hub {
@@ -229,25 +251,29 @@ func newHub(topic string, offset int64, partition int) *Hub {
 		topic:      topic,
 		partition:  partition,
 		offset:     offset,
-		reader:     create(topic, offset, partition),
+		reader:     createKafkaReader(topic, offset, partition),
 		broadcast:  make(chan []byte),
 		register:   make(chan *Client),
 		unregister: make(chan *Client),
 		clients:    make(map[*Client]bool),
 	}
 }
-func (h *Hub) read() {
+
+// Read message from kafka broker and broadcast to the clients
+
+func (h *Hub) readKafka() {
+
 	log.WithField("hub", *h).Info("start to read")
 	for h.running == true {
 
 		if h.reader == nil {
 			log.Info("waiting ... ")
 			time.Sleep(writeWait)
-			h.reader = create(h.topic, h.offset, h.partition)
+			h.reader = createKafkaReader(h.topic, h.offset, h.partition)
 		}
 		m, err := h.reader.ReadMessage(context.Background())
 		if err != nil {
-			log.WithFields(log.Fields{"error": err}).Error("read error")
+			log.WithFields(log.Fields{"error": err}).Error("websocket read")
 			h.reader.Close()
 			h.reader = nil
 		} else {
@@ -263,7 +289,7 @@ func (h *Hub) getKey() string {
 }
 
 func (h *Hub) run() {
-	go h.read()
+	go h.readKafka()
 	for h.running == true {
 		select {
 		case client := <-h.register:
@@ -273,12 +299,13 @@ func (h *Hub) run() {
 				delete(h.clients, client)
 				close(client.send)
 				log.WithField("client", client).Warn("delete client")
+				// Check if
 				if len(h.clients) < 1 {
 					log.Warn("all clients left")
-					h.reader.Close()
 					h.running = false
+					// h.reader.Close()
 					delete(hubMap, h.getKey())
-					log.WithField("hub", h).Info("clear hub")
+					log.WithField("hub", h).Info("remove hub")
 				}
 			}
 		case message := <-h.broadcast:
@@ -292,6 +319,20 @@ func (h *Hub) run() {
 			}
 		}
 	}
+	log.WithField("reader", h.reader).Warn("closing kafka reader ")
+	h.reader.Close()
+	log.WithField("hub", h).Warn("closed")
+}
+
+func destroy() {
+	log.Info(" terminating ")
+	// close all hub
+	for key, hub := range hubMap {
+		log.WithField("key", key).Warn("closing hub")
+		hub.running = false
+		// hub.reader.Close()
+	}
+	log.Fatal(" terminated ")
 }
 
 func main() {
@@ -329,6 +370,7 @@ func main() {
 	defer cancel()
 	// Doesn't block if no connections, but will otherwise wait
 	// until the timeout deadline.
+	destroy()
 	srv.Shutdown(ctx)
 	// Optionally, you could run srv.Shutdown in a goroutine and block on
 	// <-ctx.Done() if your application should wait for other services