分区和位移的参数从url输入

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")
+	// 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() {