grpc实现一个简单的聊天室

为什么80%的码农都做不了架构师？>>>   

一个简单的helloword聊天室. 先粗鲁地上个图

client side <----> server side

helloworld.proto

syntax = "proto3";

package helloword;

import "github.com/golang/protobuf/ptypes/timestamp/timestamp.proto";

service Greeter {
    rpc SayHello(stream HelloRequest) returns (stream HelloReply) {};
}

message HelloRequest {
    string message = 1;
}

message HelloReply {
    string message = 1;
    google.protobuf.Timestamp TS = 2;
    MessageType message_type = 3;

    enum MessageType{
        CONNECT_SUCCESS = 0;
        CONNECT_FAILED = 1;
        NORMAL_MESSAGE = 2;
    }
}

service:

/*
 *
 *     Author        : tuxpy
 *     Email         : q8886888@qq.com.com
 *     Create time   : 3/7/18 9:18 AM
 *     Filename      : service.go
 *     Description   :
 *
 *
 */

package main

import (
	"fmt"
	pb "grpclb/helloword"
	"log"
	"net"
	"os"
	"sync"
	"time"
	"utils"

	"github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
)

type Service struct{}

type ConnectPool struct {
	sync.Map
}

var connect_pool *ConnectPool

func (p *ConnectPool) Get(name string) pb.Greeter_SayHelloServer {
	if stream, ok := p.Load(name); ok {

		return stream.(pb.Greeter_SayHelloServer)
	} else {
		return nil
	}
}

func (p *ConnectPool) Add(name string, stream pb.Greeter_SayHelloServer) {
	p.Store(name, stream)
}

func (p *ConnectPool) Del(name string) {
	p.Delete(name)
}

func (p *ConnectPool) BroadCast(from, message string) {
	log.Printf("BroadCast from: %s, message: %s\n", from, message)
	p.Range(func(username_i, stream_i interface{}) bool {
		username := username_i.(string)
		stream := stream_i.(pb.Greeter_SayHelloServer)
		if username == from {
			return true
		} else {
			stream.Send(&pb.HelloReply{
				Message:     message,
				MessageType: pb.HelloReply_NORMAL_MESSAGE,
				TS:          &timestamp.Timestamp{Seconds: time.Now().Unix()},
			})
		}
		return true
	})
}

func (s *Service) SayHello(stream pb.Greeter_SayHelloServer) error {
	peer, _ := peer.FromContext(stream.Context())
	log.Printf("Received new connection.  %s", peer.Addr.String())

	md, _ := metadata.FromIncomingContext(stream.Context())
	username := md["name"][0]
	if connect_pool.Get(username) != nil {
		stream.Send(&pb.HelloReply{
			Message:     fmt.Sprintf("username %s already exists!", username),
			MessageType: pb.HelloReply_CONNECT_FAILED,
		})
		return nil

	} else { // 连接成功
		connect_pool.Add(username, stream)
		stream.Send(&pb.HelloReply{
			Message:     fmt.Sprintf("Connect success!"),
			MessageType: pb.HelloReply_CONNECT_SUCCESS,
		})
	}
	go func() {
		<-stream.Context().Done()
		connect_pool.Del(username)
		connect_pool.BroadCast(username, fmt.Sprintf("%s leval room", username))
	}()
	connect_pool.BroadCast(username, fmt.Sprintf("Welcome %s!", username))
	for {
		req, err := stream.Recv()
		if err != nil {
			return err
		}
		connect_pool.BroadCast(username, fmt.Sprintf("%s: %s", username, req.Message))
	}
	return nil
}

func GetListen() string {
	if len(os.Args) < 2 {
		return ":8881"
	}
	return os.Args[1]
}

func main() {
	connect_pool = &ConnectPool{}

	lis, err := net.Listen("tcp", GetListen())
	utils.CheckErrorPanic(err)
	fmt.Println("Listen on", GetListen())

	s := grpc.NewServer(grpc.RPCCompressor(grpc.NewGZIPCompressor()),
		grpc.RPCDecompressor(grpc.NewGZIPDecompressor()))
	pb.RegisterGreeterServer(s, &Service{})

	utils.CheckErrorPanic(s.Serve(lis))
}

client.go

package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"sync"
	"time"
	"utils"

	"github.com/pkg/errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	pb "grpclb/helloword"
)

var name *string = flag.String("name", "guess", "what's your name?")
var address *string = flag.String("address", "127.0.0.1:8881", "server address")
var mutex sync.Mutex

func ConsoleLog(message string) {
	mutex.Lock()
	defer mutex.Unlock()
	fmt.Printf("\n------ %s -----\n%s\n> ", time.Now(), message)
}

func Input(prompt string) string {
	fmt.Print(prompt)
	reader := bufio.NewReader(os.Stdin)
	line, _, err := reader.ReadLine()
	if err != nil {
		if err == io.EOF {
			return ""
		} else {
			panic(errors.Wrap(err, "Input"))
		}
	}
	return string(line)
}

func main() {
	flag.Parse()
	conn, err := grpc.Dial(*address, grpc.WithInsecure(),
		grpc.WithDecompressor(grpc.NewGZIPDecompressor()),
		grpc.WithCompressor(grpc.NewGZIPCompressor()),
		grpc.WithTimeout(time.Second*10))
	utils.CheckErrorPanic(errors.WithMessage(err, "grpc.Dial"))

	client := pb.NewGreeterClient(conn)
	ctx, cancel := context.WithCancel(context.Background())
	ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("name", *name))
	stream, err := client.SayHello(ctx)
	utils.CheckErrorPanic(err)

	connected := make(chan bool)
	// 监听服务端通知
	go func() {
		var (
			reply *pb.HelloReply
			err   error
		)
		for {
			reply, err = stream.Recv()
			utils.CheckErrorPanic(err)
			ConsoleLog(reply.Message)
			if reply.MessageType == pb.HelloReply_CONNECT_FAILED {
				cancel()
				break
			}
			if reply.MessageType == pb.HelloReply_CONNECT_SUCCESS {
				connected <- true
			}
		}
	}()

	go func() {
		<-connected
		var (
			line string
			err  error
		)
		for {
			line = Input("")
			if line == "exit" {
				cancel()
				break
			}
			err = stream.Send(&pb.HelloRequest{
				Message: line,
			})
			fmt.Print("> ")
			utils.CheckErrorPanic(err)
		}
	}()

	<-ctx.Done()
	fmt.Println("Bye")
}

转载于:https://my.oschina.net/tuxpy/blog/1631953