Go gRPC 发布订阅

pubsub/pubsub.proto

// protoc --go_out=plugins=grpc:. pubsub.proto

syntax = "proto3";

package proto;

message String {
    string value = 1;
}

service PubsubService {
    rpc Publish (String) returns (String);
    rpc SubscribeTopic (String) returns (stream String);
    rpc Subscribe (String) returns (stream String);
}

server/server.go

package main

import (
	"context"
	"log"
	"net"
	"strings"
	"time"

	proto "../pubsub"
	"github.com/docker/docker/pkg/pubsub"
	"google.golang.org/grpc"
)

type PubsubService struct {
	pub *pubsub.Publisher
}

func (p *PubsubService) Publish(ctx context.Context, arg *proto.String) (*proto.String, error) {
	p.pub.Publish(arg.GetValue())
	return &proto.String{}, nil
}

func (p *PubsubService) SubscribeTopic(arg *proto.String, stream proto.PubsubService_SubscribeTopicServer) error {
	ch := p.pub.SubscribeTopic(func(v interface{}) bool {
		if key, ok := v.(string); ok {
			if strings.HasPrefix(key, arg.GetValue()) {
				return true
			}
		}
		return false
	})

	for v := range ch {
		if err := stream.Send(&proto.String{Value: v.(string)}); nil != err {
			return err
		}
	}
	return nil
}

func (p *PubsubService) Subscribe(arg *proto.String, stream proto.PubsubService_SubscribeServer) error {
	ch := p.pub.Subscribe()

	for v := range ch {
		if err := stream.Send(&proto.String{Value: v.(string)}); nil != err {
			return err
		}
	}
	return nil
}

func NewPubsubService() *PubsubService {
	return &PubsubService{pub: pubsub.NewPublisher(100*time.Millisecond, 10)}
}

func main() {
	grpcServer := grpc.NewServer()
	proto.RegisterPubsubServiceServer(grpcServer, NewPubsubService())

	lis, err := net.Listen("tcp4", ":1234")
	if nil != err {
		log.Fatal(err)
	}
	grpcServer.Serve(lis)
}

client1/client1.go 用来发布信息的客户端

package main

import (
	"context"
	"log"

	proto "../pubsub"
	"google.golang.org/grpc"
)

func main() {
	conn, err := grpc.Dial("localhost:1234", grpc.WithInsecure())
	if err != nil {
		log.Fatal(err)
	}

	defer conn.Close()

	client := proto.NewPubsubServiceClient(conn)

	_, err = client.Publish(
		context.Background(), &proto.String{Value: "golang: hello Go"},
	)
	if err != nil {
		log.Fatal(err)
	}

	_, err = client.Publish(
		context.Background(), &proto.String{Value: "docker: hello Docker"},
	)
	if nil != err {
		log.Fatal(err)
	}
}

client2/client2.go 用来订阅信息的客户端

package main

import (
	"context"
	"fmt"
	"io"
	"log"

	proto "../pubsub"
	"google.golang.org/grpc"
)

func main() {
	conn, err := grpc.Dial("localhost:1234", grpc.WithInsecure())
	if nil != err {
		log.Fatal(err)
	}
	defer conn.Close()

	client := proto.NewPubsubServiceClient(conn)
	stream, err := client.Subscribe(
		context.Background(), &proto.String{Value: "golang:"},
	)
	if nil != err {
		log.Fatal(err)
	}

	go func() {
		for {
			reply, err := stream.Recv()
			if nil != err {
				if io.EOF == err {
					break
				}
				log.Fatal(err)
			}
			fmt.Println("sub1: ", reply.GetValue())
		}
	}()

	streamTopic, err := client.SubscribeTopic(
		context.Background(), &proto.String{Value: "golang:"},
	)
	if nil != err {
		log.Fatal(err)
	}

	go func() {
		for {
			reply, err := streamTopic.Recv()
			if nil != err {
				if io.EOF == err {
					break
				}
				log.Fatal(err)
			}
			fmt.Println("subTopic: ", reply.GetValue())
		}
	}()

	<-make(chan bool)
}