go + websocket实现并发下载文件及中途取消下载功能示例

Go  /  管理员 发布于 1年前   644

在go项目中使用webscoket ,实现在后台有一个下载任务，不停的下载文件，

然后实时监听下载进度，并且可以取消下载功能.

示例代码:

package main

import (
    "github.com/gin-gonic/gin"
    "github.com/google/uuid"
    "github.com/gorilla/websocket"
    "log"
    "net/http"
    "time"
)
var upgrader = websocket.Upgrader{
    ReadBufferSize:  5000,
    WriteBufferSize: 5000,
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}
type Client struct {
    uuid string
    conn *websocket.Conn
}
var ClientList []Client
var ResMessage struct {
    Type  int         `json:"type"`  // 返回的消息类型
    Data  interface{} `json:"data"`  // 返回的数据
    Total int         `json:"total"` // 在线人数
}
type Downloader struct {
    Finish int
    Total  int
    Stop   bool
}
var task *Downloader
// 使用协程模拟后台下载
func Download(d *Downloader) {
    // 并发控制
    limitChan := make(chan *Downloader, 2)
    // 执行下载
    go func() {
        for i := 0; i < d.Total; i++ {
            if d.Stop {
                log.Println("下载任务已取消")
                // 通知下载任务已取消
                limitChan <- d
                close(limitChan)
                return
            }
            time.Sleep(time.Second)
            d.Finish++
            limitChan <- d
        }
        close(limitChan)
    }()
    msg := ResMessage
    msg.Type = 1
    // 监听下载进度
    for {
        res, ok := <-limitChan
        if !ok {
            break
        }
        msg.Data = res
        msg.Total = len(ClientList)
        for _, c := range ClientList {
            err := c.conn.WriteJSON(msg)
            if err != nil {
                log.Printf("【%s】消息推送是失败", c.uuid)
                return
            }
        }
        log.Printf("下载进度：%d/%d", res.Finish, res.Total)
    }
    log.Printf("全部下载完毕")
}
// 移除client
func RemoveClient(client Client) {
    for i, c := range ClientList {
        if c.uuid == client.uuid {
            client.conn.Close()
            ClientList = append(ClientList[:i], ClientList[i+1:]...)
            break
        }
    }
    log.Printf("【%s】断开链接", client.uuid)
}

案例：

我们在后台有一个下载任务，不停的下载文件，

我们使用webscoket 实时监听下载进度，并且可以取消下载

func main() {
    router := gin.Default()
    // 开始下载
    router.GET("start", func(c *gin.Context) {
        d := Downloader{
            Finish: 0,
            Total:  30,
            Stop:   false,
        }
        task = &d
        go Download(task)
        c.JSON(200, gin.H{"message": "下载任务已启动"})
    })
    router.GET("stop", func(c *gin.Context) {
        task.Stop = true
        c.JSON(200, gin.H{"message": "取消下载"})
    })
    // 监听下载进度
    router.GET("ws", func(c *gin.Context) {
        server, err := upgrader.Upgrade(c.Writer, c.Request, nil)
        if err != nil {
            log.Printf("upgrade: %v", err)
            return
        }
        var client Client
        client.uuid = uuid.NewString()
        client.conn = server
        ClientList = append(ClientList, client)
        // 监听客户端发送的消息
        for {
            msg := ResMessage
            msgType, data, err := server.ReadMessage()
            msg.Type = msgType
            msg.Data = string(data)
            if err != nil {
                RemoveClient(client)
                return
            }
            server.WriteJSON(msg)
        }
    })
    router.Run(":3033")
}

以上是一个简单的demo , 有兴趣的可以自行测试及增加业务逻辑.