用Golang的gRPC客户端流来传输文件

在这个例子中，我们将使用gRPC客户端向服务器传输一个图像文件。我们将使用客户端流技术，所以文件将以小块的形式被传送。一旦所有的小块被传送到服务器，它将被保存。服务器将把它的唯一名称返回给客户端。

我在代码中硬编码了一些变量，但你应该使用环境变量。另外，它也是可以改进的。我试图让它尽可能的简短。

协议缓冲区

运行protoc --go_out=plugins=grpc:. --go_opt=paths=source_relative pkg/proto/*.proto 命令来生成编译文件。

// pkg/proto/upload.proto

syntax = "proto3";

package proto;

option go_package = ".;uploadpb";

service UploadService {
    rpc Upload(stream UploadRequest) returns (UploadResponse) {}
}

message UploadRequest {
    string mime = 1;
    bytes chunk = 2;
}

message UploadResponse {
    string name = 1;
}

客户端

main.go

package main

import (
	"context"
	"flag"
	"log"

	"github.com/you/transfer/internal/upload"

	"google.golang.org/grpc"
)

func main() {
	// Catch user input.
	flag.Parse()
	if flag.NArg() == 0 {
		log.Fatalln("Missing file path")
	}

	// Initialise gRPC connection.
	conn, err := grpc.Dial(":50051", grpc.WithInsecure())
	if err != nil {
		log.Fatalln(err)
	}
	defer conn.Close()

	// Start uploading the file. Error if failed, otherwise echo download URL.
	client := upload.NewClient(conn)
	name, err := client.Upload(context.Background(), flag.Arg(0))
	if err != nil {
		log.Fatalln(err)
	}
	log.Println(name)
}

upload.go

package upload

import (
	"context"
	"io"
	"os"
	"time"

	"google.golang.org/grpc"

	uploadpb "github.com/you/transfer/pkg/proto"
)

type Client struct {
	client uploadpb.UploadServiceClient
}

func NewClient(conn grpc.ClientConnInterface) Client {
	return Client{
		client: uploadpb.NewUploadServiceClient(conn),
	}
}

func (c Client) Upload(ctx context.Context, file string) (string, error) {
	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(10*time.Second))
	defer cancel()

	stream, err := c.client.Upload(ctx)
	if err != nil {
		return "", err
	}

	fil, err := os.Open(file)
	if err != nil {
		return "", err
	}

	// Maximum 1KB size per stream.
	buf := make([]byte, 1024)

	for {
		num, err := fil.Read(buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			return "", err
		}

		if err := stream.Send(&uploadpb.UploadRequest{Chunk: buf[:num]}); err != nil {
			return "", err
		}
	}

	res, err := stream.CloseAndRecv()
	if err != nil {
		return "", err
	}

	return res.GetName(), nil
}

服务器

main.go

package main

import (
	"log"
	"net"

	"github.com/you/transfer/internal/storage"
	"github.com/you/transfer/internal/upload"

	"google.golang.org/grpc"

	uploadpb "github.com/you/transfer/pkg/proto"
)

func main() {
	// Initialise TCP listener.
	lis, err := net.Listen("tcp", ":50051")
	if err != nil {
		log.Fatal(err)
	}
	defer lis.Close()

	// Bootstrap upload server.
	uplSrv := upload.NewServer(storage.New("tmp/"))

	// Bootstrap gRPC server.
	rpcSrv := grpc.NewServer()

	// Register and start gRPC server.
	uploadpb.RegisterUploadServiceServer(rpcSrv, uplSrv)
	log.Fatal(rpcSrv.Serve(lis))
}

upload.go

package upload

import (
	"io"

	"github.com/you/transfer/internal/storage"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	uploadpb "github.com/you/transfer/pkg/proto"
)

type Server struct {
	storage storage.Manager
}

func NewServer(storage storage.Manager) Server {
	return Server{
		storage: storage,
	}
}

func (s Server) Upload(stream uploadpb.UploadService_UploadServer) error {
	name := "some-unique-name.png"
	file := storage.NewFile(name)

	for {
		req, err := stream.Recv()
		if err == io.EOF {
			if err := s.storage.Store(file); err != nil {
				return status.Error(codes.Internal, err.Error())
			}

			return stream.SendAndClose(&uploadpb.UploadResponse{Name: name})
		}
		if err != nil {
			return status.Error(codes.Internal, err.Error())
		}

		if err := file.Write(req.GetChunk()); err != nil {
			return status.Error(codes.Internal, err.Error())
		}
	}
}

文件

package storage

import (
	"bytes"
)

type File struct {
	name   string
	buffer *bytes.Buffer
}

func NewFile(name string) *File {
	return &File{
		name:   name,
		buffer: &bytes.Buffer{},
	}
}

func (f *File) Write(chunk []byte) error {
	_, err := f.buffer.Write(chunk)

	return err
}

存储

package storage

import (
	"io/ioutil"
)

type Manager interface {
	Store(file *File) error
}

var _ Manager = &Storage{}

type Storage struct {
	dir string
}

func New(dir string) Storage {
	return Storage{
		dir: dir,
	}
}

func (s Storage) Store(file *File) error {
	if err := ioutil.WriteFile(s.dir+file.name, file.buffer.Bytes(), 0644); err != nil {
		return err
	}

	return nil
}

测试

// Run server first
$ go run -race cmd/server/main.go

// Upload file
$ go run -race cmd/client/main.go ~/Desktop/test.png
2021/04/13 11:30:03 some-unique-name.png

如果你检查tmp 文件夹，some-unique-name.png 文件现在应该在那里了。