每个hub单独维护自己的运行状态,没有客户端的时候,停止kafka客户端接收

This commit is contained in:
fengzhiqiang
2019-01-08 17:40:47 +08:00
parent 88f9bf7b92
commit 7a95d9350d
+56 -26
View File
@@ -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)
})