diff --git a/serv.go b/serv.go index 5f72d48..040b520 100644 --- a/serv.go +++ b/serv.go @@ -37,13 +37,14 @@ var ( newline = []byte{'\n'} 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{ + // 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, } hubMap = make(map[string]Hub) + router = mux.NewRouter() ) // Client is a middleman between the websocket connection and the hub. @@ -57,20 +58,20 @@ type Client struct { send chan []byte } -func create(topic string, offset int64) *kafka.Reader { - group, _ := uuid.NewV4() +func create(topic string, offset int64, partition int) *kafka.Reader { + group := uuid.NewV4() log.WithFields(log.Fields{ "topic": topic, "bootstrap servers": *bootstrapServers, - "partition": *partition, + "partition": partition, "group": group, "offset": offset}).Info("creating kafka client ... ") reader := kafka.NewReader(kafka.ReaderConfig{ Brokers: strings.Split(*bootstrapServers, ","), Topic: topic, - Partition: *partition, + Partition: partition, MinBytes: 10e1, MaxBytes: 10e6, // 10MB }) @@ -161,10 +162,12 @@ func (c *Client) writePump() { } // serveWs handles websocket requests from the peer. -func serveWs(w http.ResponseWriter, r *http.Request) { +func websocketHandler(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) topic := vars["topic"] - offsetString := vars["offset"] + partitionString := vars["partition"] + partition, _ := strconv.Atoi(partitionString) + offsetString := r.FormValue("offset") offset, _ := strconv.ParseInt(offsetString, 10, 64) conn, err := upgrader.Upgrade(w, r, nil) @@ -172,11 +175,14 @@ func serveWs(w http.ResponseWriter, r *http.Request) { log.Error("upgrade connection failed ", err) return } - key := topic + "|" + offsetString + key := topic + "|" + partitionString + "|" + offsetString h, ok := hubMap[key] if !ok { - log.WithFields(log.Fields{"topic": topic, "offset": offset}).Info("create new hub ") - h = *newHub(topic, offset) + log.WithFields( + log.Fields{"topic": topic, + "offset": offset, + "partition": partition}).Info("create new hub ") + h = *newHub(topic, offset, partition) go h.run() hubMap[key] = h } else { @@ -198,6 +204,8 @@ type Hub struct { topic string + partition int + offset int64 reader *kafka.Reader @@ -215,11 +223,13 @@ type Hub struct { unregister chan *Client } -func newHub(topic string, offset int64) *Hub { +func newHub(topic string, offset int64, partition int) *Hub { return &Hub{ running: true, topic: topic, + partition: partition, offset: offset, + reader: create(topic, offset, partition), broadcast: make(chan []byte), register: make(chan *Client), unregister: make(chan *Client), @@ -233,7 +243,7 @@ func (h *Hub) read() { if h.reader == nil { log.Info("waiting ... ") time.Sleep(writeWait) - h.reader = create(h.topic, h.offset) + h.reader = create(h.topic, h.offset, h.partition) } m, err := h.reader.ReadMessage(context.Background()) if err != nil { @@ -249,7 +259,7 @@ func (h *Hub) read() { } func (h *Hub) getKey() string { - return h.topic + "|" + strconv.FormatInt(h.offset, 10) + return h.topic + "|" + strconv.Itoa(h.partition) + "|" + strconv.FormatInt(h.offset, 10) } func (h *Hub) run() { @@ -286,8 +296,7 @@ func (h *Hub) run() { func main() { flag.Parse() - r := mux.NewRouter() - loggedRouter := handlers.LoggingHandler(os.Stdout, r) + loggedRouter := handlers.LoggingHandler(os.Stdout, router) srv := &http.Server{ Addr: *addr, // Good practice to set timeouts to avoid Slowloris attacks. @@ -299,9 +308,10 @@ func main() { // hub := newHub() // go hub.run() - r.HandleFunc("/ws/{topic}/{offset}", func(w http.ResponseWriter, r *http.Request) { - serveWs(w, r) - }) + // r.HandleFunc("/ws/{topic}/{offset}", func(w http.ResponseWriter, r *http.Request) { + // serveWs(w, r) + // }) + router.Path("/ws/{topic}/{partition}").Queries("offset", "{offset}").HandlerFunc(websocketHandler).Name("web-socket") // Run our server in a goroutine so that it doesn't block. go func() {