ms.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. package main
  2. import (
  3. "io"
  4. "math/rand"
  5. "net/http"
  6. "os"
  7. "strconv"
  8. "time"
  9. "github.com/gin-gonic/gin"
  10. "github.com/gorilla/websocket"
  11. )
  12. func main() {
  13. port := "8080"
  14. if os.Getenv("PORT") != "" {
  15. port = os.Getenv("PORT")
  16. }
  17. serverName := "localhost"
  18. if os.Getenv("SERVER_NAME") != "" {
  19. serverName = os.Getenv("SERVER_NAME")
  20. }
  21. secureProtocol := ""
  22. if os.Getenv("SECURE_PROTOCOL") == "1" {
  23. secureProtocol = "s"
  24. }
  25. rand.Seed(time.Now().Unix())
  26. router := gin.Default()
  27. router.LoadHTMLFiles("index.html")
  28. router.GET("/", func(c *gin.Context) {
  29. c.HTML(http.StatusOK, "index.html", gin.H{
  30. "websocket": "ws" + secureProtocol + "://" + serverName + ":" + port + "/ws",
  31. "sse": "http" + secureProtocol + "://" + serverName + ":" + port + "/sse",
  32. })
  33. })
  34. router.GET("/ws", websocketHandler(websocket.Upgrader{
  35. Subprotocols: []string{"protocolOne"},
  36. }))
  37. router.GET("/sse", serverSentEventHandler())
  38. router.Run(":" + port)
  39. }
  40. func websocketHandler(upgrader websocket.Upgrader) func(c *gin.Context) {
  41. return func(c *gin.Context) {
  42. conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
  43. if err != nil {
  44. c.AbortWithError(http.StatusInternalServerError, err)
  45. return
  46. }
  47. t := task{}
  48. t.start()
  49. for {
  50. progress, ok := <-t.progress
  51. if !ok {
  52. break
  53. }
  54. if err := conn.WriteMessage(1, []byte(strconv.Itoa(progress))); err != nil {
  55. c.AbortWithError(http.StatusInternalServerError, err)
  56. return
  57. }
  58. }
  59. }
  60. }
  61. func serverSentEventHandler() func(c *gin.Context) {
  62. return func(c *gin.Context) {
  63. t := task{}
  64. t.start()
  65. c.Stream(func(w io.Writer) bool {
  66. progress, ok := <-t.progress
  67. if !ok {
  68. return false
  69. }
  70. c.SSEvent("", strconv.Itoa(progress))
  71. return true
  72. })
  73. }
  74. }
  75. type task struct {
  76. progress chan int
  77. }
  78. func (t *task) start() {
  79. t.progress = make(chan int)
  80. go func() {
  81. p := 0
  82. t.progress <- p
  83. for p < 100 {
  84. if rand.Intn(100) > 50 {
  85. p++
  86. t.progress <- p
  87. }
  88. time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
  89. }
  90. close(t.progress)
  91. }()
  92. }