// Demo will serve an HTML page that demonstrates how to use the 'stream'
// endpoint.
func (s *StreamService) Demo(w http.ResponseWriter, r *http.Request) {
vals := struct {
Port int
StreamID int64
}{
s.port,
web.GetInt64Var(r, "stream_id"),
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
err := demoTempl.Execute(w, &vals)
if err != nil {
server.Log.Error("template error ", err)
http.Error(w, "problems loading HTML", http.StatusInternalServerError)
}
}
func TestGetInt64Var(t *testing.T) {
tests := []struct {
givenURL string
givenRoute string
want int64
}{
{
"/blah/123",
"/blah/{key}",
123,
},
{
"/blah/adsf",
"/blah/{key}",
0,
},
{
"/blah?key=123",
"/blah",
123,
},
{
"/blah?key=abc",
"/blah",
0,
},
}
for _, test := range tests {
route := mux.NewRouter()
route.HandleFunc(test.givenRoute, func(w http.ResponseWriter, r *http.Request) {
web.SetRouteVars(r, mux.Vars(r))
got := web.GetInt64Var(r, "key")
if got != test.want {
t.Errorf("got int of %#v, expected %#v", got, test.want)
}
})
r, _ := http.NewRequest("GET", test.givenURL, nil)
route.ServeHTTP(httptest.NewRecorder(), r)
}
}
// Stream will init a new pubsub.KafkaPublisher and pubsub.KafkaSubscriber
// then upgrade the current request to a websocket connection. Any messages
// consumed from Kafka will be published to the web socket and vice versa.
func (s *StreamService) Stream(w http.ResponseWriter, r *http.Request) {
cfg := *s.cfg
cfg.Topic = topicName(web.GetInt64Var(r, "stream_id"))
server.LogWithFields(r).WithField("topic", cfg.Topic).Info("new stream req")
sub, err := pubsub.NewKafkaSubscriber(&cfg, zeroOffset, discardOffset)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to create sub")
http.Error(w, "unable to create subscriber: "+err.Error(), http.StatusBadRequest)
return
}
defer func() {
if err := sub.Stop(); err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to stop sub")
}
}()
var pub *pubsub.KafkaPublisher
pub, err = pubsub.NewKafkaPublisher(&cfg)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to create pub")
http.Error(w, "unable to create publisher: "+err.Error(), http.StatusBadRequest)
return
}
defer func() {
if err := pub.Stop(); err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to stop pub")
}
}()
var ws *websocket.Conn
ws, err = upgrader.Upgrade(w, r, nil)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to create websocket")
http.Error(w, "unable to create websocket: "+err.Error(), http.StatusBadRequest)
return
}
defer func() {
if err := ws.Close(); err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to close ws")
}
}()
// start consumer, emit to ws
noopTicker := time.NewTicker(time.Second * 5)
subscriberDone := make(chan bool, 1)
stopSubscriber := make(chan bool, 1)
go func() {
defer func() { subscriberDone <- true }()
var (
payload []byte
msgs = sub.Start()
)
for {
select {
case msg := <-msgs:
payload = msg.Message()
err = ws.SetWriteDeadline(time.Now().Add(time.Second * 30))
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to write deadline")
}
err = ws.WriteMessage(websocket.TextMessage, payload)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to write ws message")
}
err = msg.Done()
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to mark gizmo message as done")
}
case <-noopTicker.C:
err = ws.SetWriteDeadline(time.Now().Add(time.Second * 30))
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to write deadline")
}
if err := ws.WriteMessage(websocket.PingMessage, []byte("ping")); err != nil {
server.LogWithFields(r).WithField("error", err).Error("error writing ws message")
return
}
case <-stopSubscriber:
return
}
}
}()
// start producer, emit to kafka
producerDone := make(chan bool, 1)
go func() {
defer func() {
producerDone <- true
stopSubscriber <- true
server.LogWithFields(r).WithField("topic", cfg.Topic).Info("closing stream req")
}()
var (
messageType int
payload []byte
err error
read io.Reader
)
for {
messageType, read, err = ws.NextReader()
if err != nil {
if err != io.EOF {
server.LogWithFields(r).WithField("error", err).Error("error reading message")
}
return
}
switch messageType {
case websocket.TextMessage:
payload, err = ioutil.ReadAll(read)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to read payload")
return
}
err = pub.PublishRaw(cfg.Topic, payload)
if err != nil {
server.LogWithFields(r).WithField("error", err).Error("unable to publish payload")
return
}
case websocket.PingMessage, websocket.PongMessage, websocket.BinaryMessage:
server.LogWithFields(r).Info("discarding message type: ", messageType)
case websocket.CloseMessage:
server.LogWithFields(r).Info("closing websocket")
return
}
}
}()
<-subscriberDone
<-producerDone
noopTicker.Stop()
server.LogWithFields(r).WithField("topic", cfg.Topic).Info("leaving stream req")
}