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") } secureProtocol := "" if os.Getenv("SECURE_PROTOCOL") == "1" { secureProtocol = "s" } 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" + secureProtocol + "://" + serverName + ":" + port + "/ws", "sse": "http" + secureProtocol + "://" + 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) }() }