kafka-proxy/serv.go

package main

import (
	"bytes"
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	uuid "github.com/satori/go.uuid"
	kafka "github.com/segmentio/kafka-go"
	log "github.com/sirupsen/logrus"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	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
)

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:  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

	// The websocket connection.
	conn *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte
}

func createKafkaReader(topic string, offset int64, partition int) *kafka.Reader {

	// Use uuid for kafka group id
	group := uuid.NewV4()

	log.WithFields(log.Fields{
		"topic":             topic,
		"bootstrap servers": *bootstrapServers,
		"partition":         partition,
		"group":             group,
		"offset":            offset}).Info("creating kafka client ... ")

	reader := kafka.NewReader(kafka.ReaderConfig{
		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
}

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

	// 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 })

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.WithError(err).Error("read message error")
			}
			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.
//
// A goroutine running writePump is started for each connection. The
// 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()
		c.conn.Close()
		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
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				log.WithError(err).Error("can't get next wirter")
				return
			}
			w.Write(message)

			// Add queued chat messages to the current websocket message.
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-c.send)
			}

			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
			}
		}
	}
}

// serveWs handles websocket requests from the peer.
func websocketHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	topic := vars["topic"]
	partitionString := vars["partition"]
	partition, _ := strconv.Atoi(partitionString)
	offsetString := r.FormValue("offset")
	offset, _ := strconv.ParseInt(offsetString, 10, 64)

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error("upgrade connection failed ", err)
		return
	}
	key := topic + "|" + partitionString + "|" + offsetString
	h, ok := hubMap[key]
	if !ok {
		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 {
		log.Info("join hub")
	}
	client := &Client{hub: &h, conn: conn, send: make(chan []byte, 256)}
	client.hub.register <- client

	// Allow collection of memory referenced by the caller by doing all work in
	// new goroutines.
	go client.writePump()
	go client.startClient()
}

func newHub(topic string, offset int64, partition int) *Hub {
	return &Hub{
		running:    true,
		topic:      topic,
		partition:  partition,
		offset:     offset,
		reader:     createKafkaReader(topic, offset, partition),
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[*Client]bool),
	}
}

// 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 = createKafkaReader(h.topic, h.offset, h.partition)
		}
		m, err := h.reader.ReadMessage(context.Background())
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Error("websocket read")
			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
		}

	}
}

func (h *Hub) getKey() string {
	return h.topic + "|" + strconv.Itoa(h.partition) + "|" + strconv.FormatInt(h.offset, 10)
}

func (h *Hub) run() {
	go h.readKafka()
	for h.running == true {
		select {
		case client := <-h.register:
			h.clients[client] = true
		case client := <-h.unregister:
			if _, ok := h.clients[client]; ok {
				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.running = false
					// h.reader.Close()
					delete(hubMap, h.getKey())
					log.WithField("hub", h).Info("remove hub")
				}
			}
		case message := <-h.broadcast:
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
		}
	}
	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() {
	flag.Parse()
	loggedRouter := handlers.LoggingHandler(os.Stdout, router)
	srv := &http.Server{
		Addr: *addr,
		// Good practice to set timeouts to avoid Slowloris attacks.
		WriteTimeout: time.Second * 15,
		ReadTimeout:  time.Second * 15,
		IdleTimeout:  time.Second * 60,
		Handler:      loggedRouter, // Pass our instance of gorilla/mux in.
	}

	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() {
		log.WithFields(log.Fields{"address": *addr}).Info("starting server")
		if err := srv.ListenAndServe(); err != nil {
			log.WithFields(log.Fields{"error": err}).Error("error")
		}
	}()

	c := make(chan os.Signal, 1)
	// We'll accept graceful shutdowns when quit via SIGINT (Ctrl+C)
	// SIGKILL, SIGQUIT or SIGTERM (Ctrl+/) will not be caught.
	signal.Notify(c, os.Interrupt)

	// Block until we receive our signal.
	<-c

	// Create a deadline to wait for.
	ctx, cancel := context.WithTimeout(context.Background(), *wait)
	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
	// to finalize based on context cancellation.
	log.Info("shutting down")
	os.Exit(0)
}