diff --git a/serv.go b/serv.go index 48441c7..16e0081 100644 --- a/serv.go +++ b/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