kafka-proxy/serv.go

package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"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 = 60 * 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'}
	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:  1024,
		WriteBufferSize: 1024,
	}
	hubMap = make(map[string]Hub)
)

// 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 create(topic string, offset int64) *kafka.Reader {
	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, ","),
		Topic:     topic,
		Partition: *partition,
		MinBytes:  10e1,
		MaxBytes:  10e6, // 10MB
	})
	reader.SetOffset(offset)

	return reader
}

// readPump pumps messages from the websocket connection to the hub.
//
// The application runs readPump in a per-connection goroutine. The application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
		log.WithField("client", c).Warn("exit")
	}()
	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 })

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

}

// 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() {
	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 {
				// The hub closed the channel.
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				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 {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// serveWs handles websocket requests from the peer.
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
	}
	key := topic + "|" + offsetString
	h, ok := hubMap[key]
	if !ok {
		log.WithFields(log.Fields{"topic": topic, "offset": offset}).Info("create new hub ")
		h = *newHub(topic, offset)
		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.readPump()
}

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

func newHub(topic string, offset int64) *Hub {
	return &Hub{
		running:    true,
		topic:      topic,
		offset:     offset,
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[*Client]bool),
	}
}
func (h *Hub) read() {
	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
		}

	}
}

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

func (h *Hub) run() {
	go h.read()
	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")
				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 {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
		}
	}
}

func main() {
	flag.Parse()
	r := mux.NewRouter()
	loggedRouter := handlers.LoggingHandler(os.Stdout, r)
	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.
	}

	// hub := newHub()
	// go hub.run()
	r.HandleFunc("/ws/{topic}/{offset}", func(w http.ResponseWriter, r *http.Request) {
		serveWs(w, r)
	})

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