From 7a95d9350dafb7e2f5393a4e84ae579a767b1d2b Mon Sep 17 00:00:00 2001 From: fengzhiqiang Date: Tue, 8 Jan 2019 17:40:47 +0800 Subject: [PATCH] =?UTF-8?q?=E6=AF=8F=E4=B8=AAhub=E5=8D=95=E7=8B=AC?= =?UTF-8?q?=E7=BB=B4=E6=8A=A4=E8=87=AA=E5=B7=B1=E7=9A=84=E8=BF=90=E8=A1=8C?= =?UTF-8?q?=E7=8A=B6=E6=80=81=EF=BC=8C=E6=B2=A1=E6=9C=89=E5=AE=A2=E6=88=B7?= =?UTF-8?q?=E7=AB=AF=E7=9A=84=E6=97=B6=E5=80=99=EF=BC=8C=E5=81=9C=E6=AD=A2?= =?UTF-8?q?kafka=E5=AE=A2=E6=88=B7=E7=AB=AF=E6=8E=A5=E6=94=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- serv.go | 82 +++++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 56 insertions(+), 26 deletions(-) diff --git a/serv.go b/serv.go index 564bf7a..55da298 100644 --- a/serv.go +++ b/serv.go @@ -6,6 +6,7 @@ import ( "net/http" "os" "os/signal" + "strconv" "strings" "syscall" "time" @@ -43,7 +44,6 @@ var ( WriteBufferSize: 1024, } hubMap = make(map[string]Hub) - run = true ) // Client is a middleman between the websocket connection and the hub. @@ -57,14 +57,15 @@ type Client struct { send chan []byte } -func create(topic string) *kafka.Reader { +func create(topic string, offset int64) *kafka.Reader { group, _ := uuid.NewV4() log.WithFields(log.Fields{ "topic": topic, "bootstrap servers": *bootstrapServers, "partition": *partition, - "group": group}).Info("creating kafka client ... ") + "group": group, + "offset": offset}).Info("creating kafka client ... ") reader := kafka.NewReader(kafka.ReaderConfig{ Brokers: strings.Split(*bootstrapServers, ","), @@ -73,7 +74,7 @@ func create(topic string) *kafka.Reader { MinBytes: 10e1, MaxBytes: 10e6, // 10MB }) - reader.SetOffset(-1) //latest offset + reader.SetOffset(offset) return reader } @@ -95,22 +96,17 @@ func (c *Client) readPump() { 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 + } log.Fatal("break ") - // default: - // for { - // m, err := c.hub.reader.ReadMessage(context.Background()) - // if err != nil { - // log.WithFields(log.Fields{"error": err}).Error("read error") - // } - // log.WithFields(log.Fields{"offset": m.Offset, "message": string(m.Value)}).Info("got message") - // c.hub.broadcast <- m.Value - // } } } @@ -126,6 +122,8 @@ func (c *Client) writePump() { defer func() { ticker.Stop() c.conn.Close() + c.hub.unregister <- c + log.WithField("client", c).Warn("exit") }() for { select { @@ -166,18 +164,21 @@ func (c *Client) writePump() { func serveWs(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) topic := vars["topic"] + offsetString := vars["offset"] + offset, _ := strconv.ParseInt(offsetString, 10, 64) conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Error("upgrade connection failed ", err) return } - h, ok := hubMap[topic] + key := topic + "|" + offsetString + h, ok := hubMap[key] if !ok { - log.WithField("topic", topic).Info("create new hub ") - h = *newHub(topic) + log.WithFields(log.Fields{"topic": topic, "offset": offset}).Info("create new hub ") + h = *newHub(topic, offset) go h.run() - hubMap[topic] = h + hubMap[key] = h } else { log.Info("join hub") } @@ -193,6 +194,12 @@ func serveWs(w http.ResponseWriter, r *http.Request) { // Hub maintains the set of active clients and broadcasts messages to the // clients. type Hub struct { + running bool + + topic string + + offset int64 + reader *kafka.Reader // Registered clients. @@ -208,9 +215,11 @@ type Hub struct { unregister chan *Client } -func newHub(topic string) *Hub { +func newHub(topic string, offset int64) *Hub { return &Hub{ - reader: create(topic), + running: true, + topic: topic, + offset: offset, broadcast: make(chan []byte), register: make(chan *Client), unregister: make(chan *Client), @@ -218,20 +227,33 @@ func newHub(topic string) *Hub { } } func (h *Hub) read() { - log.WithField("reader", h.reader).Info("start to read") - for run == true { + log.WithField("hub", *h).Info("start to read") + for h.running == true { + + if h.reader == nil { + h.reader = create(h.topic, h.offset) + time.Sleep(5 * time.Second) + } m, err := h.reader.ReadMessage(context.Background()) if err != nil { log.WithFields(log.Fields{"error": err}).Error("read error") + h.reader.Close() + h.reader = nil + } else { + log.WithFields(log.Fields{"offset": m.Offset, "message": string(m.Value)}).Info("got message") + h.broadcast <- m.Value } - log.WithFields(log.Fields{"offset": m.Offset, "message": string(m.Value)}).Info("got message") - h.broadcast <- m.Value + } } +func (h *Hub) getKey() string { + return h.topic + "|" + strconv.FormatInt(h.offset, 10) +} + func (h *Hub) run() { go h.read() - for run == true { + for h.running == true { select { case client := <-h.register: h.clients[client] = true @@ -239,6 +261,14 @@ func (h *Hub) run() { if _, ok := h.clients[client]; ok { delete(h.clients, client) close(client.send) + log.WithField("client", client).Warn("delete client") + if len(h.clients) < 1 { + log.Warn("all clients left") + h.reader.Close() + h.running = false + delete(hubMap, h.getKey()) + log.WithField("hub", h).Info("clear hub") + } } case message := <-h.broadcast: for client := range h.clients { @@ -268,7 +298,7 @@ func main() { // hub := newHub() // go hub.run() - r.HandleFunc("/ws/{topic}", func(w http.ResponseWriter, r *http.Request) { + r.HandleFunc("/ws/{topic}/{offset}", func(w http.ResponseWriter, r *http.Request) { serveWs(w, r) })