ms/ms.go

package main

import (
	"io"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

func main() {
	port := "8080"
	if os.Getenv("PORT") != "" {
		port = os.Getenv("PORT")
	}
	serverName := "localhost"
	if os.Getenv("SERVER_NAME") != "" {
		serverName = os.Getenv("SERVER_NAME")
	}

	rand.Seed(time.Now().Unix())

	router := gin.Default()

	router.LoadHTMLFiles("index.html")

	router.GET("/", func(c *gin.Context) {
		c.HTML(http.StatusOK, "index.html", gin.H{
			"websocket": "ws://" + serverName + ":" + port + "/ws",
			"sse":       "http://" + serverName + ":" + port + "/sse",
		})
	})
	router.GET("/ws", websocketHandler(websocket.Upgrader{
		Subprotocols: []string{"protocolOne"},
	}))
	router.GET("/sse", serverSentEventHandler())

	router.Run(":" + port)
}

func websocketHandler(upgrader websocket.Upgrader) func(c *gin.Context) {
	return func(c *gin.Context) {
		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		t := task{}
		t.start()
		for {
			progress, ok := <-t.progress
			if !ok {
				break
			}

			if err := conn.WriteMessage(1, []byte(strconv.Itoa(progress))); err != nil {
				c.AbortWithError(http.StatusInternalServerError, err)
				return
			}
		}
	}
}

func serverSentEventHandler() func(c *gin.Context) {
	return func(c *gin.Context) {
		t := task{}
		t.start()

		c.Stream(func(w io.Writer) bool {
			progress, ok := <-t.progress
			if !ok {
				return false
			}

			c.SSEvent("", strconv.Itoa(progress))
			return true
		})
	}
}

type task struct {
	progress chan int
}

func (t *task) start() {
	t.progress = make(chan int)
	go func() {
		p := 0
		t.progress <- p
		for p < 100 {
			if rand.Intn(100) > 50 {
				p++
				t.progress <- p
			}
			time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
		}
		close(t.progress)
	}()
}