[ go-zero的rpc服务使用流stream传输文件 ]


一.zero-api

.api文件(次要,对流没影响)

type MusicStreamDataReq {
    Output string `path:"output,default=json"`
    Scene  string `path:"scene,optional"`
    Path   string `path:"path,optional"`
}

type MusicStreamDateResp {
    Path   string `json:"path"`
    Md5    string `json:"md5"`
    Domain string `json:"domain"`
    Size   int64  `json:"size"`
}

service musicapi-api {
    @handler MusicStreamData
    post /music/stream(MusicStreamDataReq) returns (MusicStreamDateResp)
}

1. zpi-handler.go

package handler

import (
    "net/http"

    "go-weixin/app/music/musicapi/internal/logic"
    "go-weixin/app/music/musicapi/internal/svc"
    "go-weixin/app/music/musicapi/internal/types"

    "github.com/zeromicro/go-zero/rest/httpx"
    "github.com/zeromicro/x/errors"
    xhttp "github.com/zeromicro/x/http"
)

func MusicStreamDataHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        var req types.MusicStreamDataReq
        if err := httpx.Parse(r, &req); err != nil {
            httpx.ErrorCtx(r.Context(), w, errors.New(2000, err.Error()))
            return
        }
        files, fileHeader, err := r.FormFile("file")
        if err != nil {
            httpx.ErrorCtx(r.Context(), w, errors.New(2000, err.Error()))
            return
        }
        l := logic.NewMusicStreamDataLogic(r.Context(), svcCtx)
        resp, err := l.MusicStreamData(&req, files, fileHeader)
        if err != nil {
            xhttp.JsonBaseResponseCtx(r.Context(), w, err)
        } else {
            xhttp.JsonBaseResponseCtx(r.Context(), w, resp)
        }
    }
}

2. api-logic.go

package logic

import (
    "bytes"
    "context"
    "io"
    "mime/multipart"
    "time"

    "go-weixin/app/music/musicapi/internal/svc"
    "go-weixin/app/music/musicapi/internal/types"
    "go-weixin/app/music/musicrpc/pb/proto/music"

    "github.com/zeromicro/go-zero/core/logx"
    "github.com/zeromicro/x/errors"
)

type MusicStreamDataLogic struct {
    logx.Logger
    ctx    context.Context
    svcCtx *svc.ServiceContext
}

func NewMusicStreamDataLogic(ctx context.Context, svcCtx *svc.ServiceContext) *MusicStreamDataLogic {
    return &MusicStreamDataLogic{
        Logger: logx.WithContext(ctx),
        ctx:    ctx,
        svcCtx: svcCtx,
    }
}

func (l *MusicStreamDataLogic) MusicStreamData(req *types.MusicStreamDataReq, file multipart.File, fileHeader *multipart.FileHeader) (resp *types.MusicStreamDateResp, err error) {
    ctx, canelFunc := context.WithDeadline(context.Background(), time.Now().Add(10*time.Second))
    defer canelFunc()
    byteData := make([]byte, fileHeader.Size)
    file.Read(byteData)
    defer file.Close()
    stream, err := l.svcCtx.MusicRpc.MusicStreamData(ctx)
    reader := bytes.NewReader(byteData)
    buf := make([]byte, 1024)
    for {
        num, err := reader.Read(buf)
        if err == io.EOF {
            break
        }
        if err != nil {
            return nil, errors.New(7000, "读取文件异常: "+err.Error())
        }
        err = stream.Send(&music.MusicStreamDataReq{
            Filename: fileHeader.Filename,
            Size:     fileHeader.Size,
            FileByte: buf[:num],
        })
        if err != nil {
            return nil, errors.New(7000, "文件rpc流传输异常"+err.Error())
        }
    }
    res, err := stream.CloseAndRecv()
    if err != nil {
        return nil, errors.New(7000, "rpc返回异常"+err.Error())
    }
    resp = new(types.MusicStreamDateResp)
    resp.Domain = res.Domain
    resp.Md5 = res.Md5
    resp.Path = res.Path
    resp.Size = res.Size
    return
}

二. zero-rpc

proto文件

// 声明 proto 语法版本,固定值
syntax = "proto3";

package music;

option go_package = "proto/music";

message MusicStreamDataReq {
  bytes fileByte = 1;
  string filename = 2;
  int64 size = 3;
}

message MusicStreamDataResp {
  string path = 1;
  string md5 = 2;
  string domain =3;
  int64 size = 4;
}

service music {
  rpc MusicStreamData(stream MusicStreamDataReq) returns (MusicStreamDataResp);
}

1. rpc-logic.go

package logic

import (
    "context"
    "encoding/json"
    "io"
    "time"

    "go-weixin/app/music/musicrpc/internal/svc"
    "go-weixin/app/music/musicrpc/pb/proto/music"
    "go-weixin/common/utils"

    "github.com/zeromicro/go-zero/core/logx"
    "github.com/zeromicro/x/errors"
)

type MusicStreamDataLogic struct {
    ctx    context.Context
    svcCtx *svc.ServiceContext
    logx.Logger
}

func NewMusicStreamDataLogic(ctx context.Context, svcCtx *svc.ServiceContext) *MusicStreamDataLogic {
    return &MusicStreamDataLogic{
        ctx:    ctx,
        svcCtx: svcCtx,
        Logger: logx.WithContext(ctx),
    }
}

func (l *MusicStreamDataLogic) MusicStreamData(stream music.Music_MusicStreamDataServer) error {
    var fileByte []byte
    req := &music.MusicStreamDataReq{}
    res := &music.MusicStreamDataResp{}
    for {
        da, err := stream.Recv()

        if (da != nil) && (len(da.FileByte) > 0) {
            fileByte = append(fileByte, da.GetFileByte()...)
        }
        if err == io.EOF {
            fileName := ""
            var fileSize int64
            if req.Filename == "" {
                fileName = time.Now().Local().Format("20060102150405")
            } else {
                fileName = req.Filename
            }
            if req.Size == 0 {
                fileSize = int64(len(fileByte))
            } else {
                fileSize = req.Size
            }
            resByte, err := utils.GoFastDfsClient(fileByte, fileName, fileSize)
            if err != nil {
                return errors.New(8000, "文件上传异常:"+err.Error())
            }
            if err = json.Unmarshal(resByte, res); err != nil {
                return errors.New(8000, "json序列化异常:"+err.Error())
            }
            stream.SendAndClose(res)
            return nil
        }
        if err != nil {
            stream.SendAndClose(res)
            return errors.New(8000, "rpc流返回异常:"+err.Error())
        }
        req.Filename = da.GetFilename()
        req.Size = da.GetSize()

    }
}

utils

1. upload.go

package utils

import (
    "bytes"
    "fmt"
    "io"
    "mime/multipart"
    "net/http"
)

func GoFastDfsClient(fileByte []byte, fileName string, fileSize int64) ([]byte, error) {
    url := "http://192.168.50.205:8080/group1/upload"
    bodyBuf := bytes.NewBufferString("")
    bodyWrite := multipart.NewWriter(bodyBuf)
    _, err := bodyWrite.CreateFormFile("file", fileName)
    if err != nil {
        return nil, err
    }

    boundary := bodyWrite.Boundary()
    close_buf := bytes.NewBufferString(fmt.Sprintf("\r\n--%s--\r\n", boundary))
    file := bytes.NewReader(fileByte)
    request_render := io.MultiReader(bodyBuf, file, close_buf)

    httpReq, err := http.NewRequest(http.MethodPost, url, request_render)
    if err != nil {
        return nil, err
    }
    query := httpReq.URL.Query()
    query.Add("output", "json")
    query.Add("scene", "")
    query.Add("path", "")
    httpReq.URL.RawQuery = query.Encode()
    httpReq.Header.Add("Content-Type", "multipart/form-data; boundary="+boundary)
    httpReq.ContentLength = fileSize + int64(bodyBuf.Len()) + int64(close_buf.Len())
    httpRes, err := http.DefaultClient.Do(httpReq)
    if err != nil {
        return nil, err
    }
    defer httpRes.Body.Close()
    resBody, err := io.ReadAll(httpRes.Body)
    if err != nil {
        return nil, err
    }

    return resBody, nil
}