mongo_socks.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	_ "github.com/joho/godotenv/autoload"
	log "github.com/sirupsen/logrus"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"

	"github.ibm.com/CIOCloud/lumberjack/db"
	"github.ibm.com/CIOCloud/lumberjack/util"
)

const host = "localhost:8081"

func helloHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "html/index.html")
}

func websocketHandler(mg *mongo.Client, w http.ResponseWriter, r *http.Request) {
	build := r.URL.Query().Get("build")
	if build == "" {
		http.Error(w, "invalid build specified", http.StatusBadRequest)
		return
	}

	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Errorln(err.Error())
		return
	}
	defer conn.Close()
	log.Infoln("client connected:", r.RemoteAddr)

	ctx, cancel := context.WithCancel(context.Background())

	logs := make(chan bson.M)

	go func() {
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Infoln("client disconnected:", r.RemoteAddr)
				cancel()
				return
			}
			log.Infoln("message from client:", message)
		}
	}()

	go func() {
		col := mg.Database("build_logs").Collection(build)
		stream, err := col.Watch(ctx, mongo.Pipeline{
			bson.D{{
				"$match",
				bson.D{{"operationType", "insert"}},
			}},
		})
		defer stream.Close(ctx)
		if err != nil {
			log.Errorln(err)
			cancel()
			return
		}

		for stream.Next(ctx) {
			var event bson.M
			err = stream.Decode(&event)
			if err != nil {
				log.Errorln(err)
				cancel()
				return
			}

			logs <- event
		}
	}()

	for {
		logline := <-logs
		fmt.Println(logline)
		err = conn.WriteJSON(logline)
		if err != nil {
			log.Infoln("client disconnected:", r.RemoteAddr)
			cancel()
			return
		}
	}

	// setup websocket
	// start listening from mongo into 'logs' channel
	// get all logs before "now" -> prev logs
	// send prev logs
	// work thru channel to send curr logs
}

func main() {
	ctx := context.Background()
	mg, err := db.Connect(ctx)
	defer mg.Disconnect(ctx)
	if err != nil {
		log.Fatalln("unable to establish mongo connection:", err)
	}

	timeoutCtx, _ := context.WithTimeout(ctx, 15*time.Second)
	err = mg.Ping(timeoutCtx, nil)
	if err != nil {
		log.Fatalln("unable to ping mongo:", err)
	}

	http.HandleFunc("/", helloHandler)
	http.Handle("/ws", &util.MongoHandler{DB: mg, Handler: websocketHandler})

	log.Infoln("Starting Server on:", host)
	log.Fatalln(http.ListenAndServe(host, nil))
}