sonet/pkg/protocol/deliver/deliver.go

package deliver

import (
	"context"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/proto"
	"reflect"
	"sonet/api/gen/postal"
	"sonet/pkg/grpc/balancer"
	"sonet/pkg/grpc/discovery"
	"sonet/pkg/utils/logger"
	"time"
)

type Status int16

const (
	StatusSuccess         Status = 1
	StatusError           Status = 2
	StatusReceiverOffline Status = 10
)

// Deliver n包,通知消息投递
type Deliver struct {
	svcName string
	postal  postal.PostalClient
}

func NewDeliver(msgInServiceName string) *Deliver {
	return &Deliver{
		svcName: msgInServiceName,
	}
}

func (d *Deliver) InitWithResolver(ctx context.Context, resolver discovery.GrpcResolver, opts ...grpc.DialOption) (err error) {
	balancer.InitConsistentHashBuilder()
	rb, err := resolver.Resolver()
	if err != nil {
		return
	}

	postalUrl := resolver.DialUrl(postal.Postal_ServiceDesc.ServiceName)
	var options []grpc.DialOption
	// consistent hash lb
	options = append(options, grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingPolicy":"%s"}`, balancer.ConsistentHash)))
	options = append(options, grpc.WithResolvers(rb))
	options = append(options, grpc.WithTransportCredentials(insecure.NewCredentials()))
	options = append(options, opts...)
	conn, err := grpc.DialContext(ctx, postalUrl, options...)
	if err != nil {
		return
	}
	d.postal = postal.NewPostalClient(conn)
	return
}

func (d *Deliver) InitWithAddr(postalAddr string, opts ...grpc.DialOption) (err error) {
	// Conn *grpc.ClientConn
	var options []grpc.DialOption
	options = append(options, grpc.WithTransportCredentials(insecure.NewCredentials()))
	options = append(options, opts...)
	conn, err := grpc.Dial(postalAddr, options...)
	if err != nil {
		return
	}
	d.postal = postal.NewPostalClient(conn)
	return
}

func (d *Deliver) Deliver(ctx context.Context, msg proto.Message, receiver string, options ...Option) (Status, error) {
	return d.deliver0(ctx, msg, []string{receiver}, options...)
}

func (d *Deliver) DeliverBatch(ctx context.Context, msg proto.Message, receivers []string, options ...Option) (Status, error) {
	return d.deliver0(ctx, msg, receivers, options...)
}

func protoMessage2Deliver(svcName string, msg proto.Message) (message *postal.Message, err error) {
	body, err := proto.Marshal(msg)
	if err != nil {
		return
	}
	msgName := reflect.TypeOf(msg).Elem().Name()
	message = &postal.Message{
		Time: time.Now().UnixMilli(),
		Svc:  svcName,
		Msg:  msgName,
		Body: body,
	}
	return
}

func (d *Deliver) deliver0(ctx context.Context, msg proto.Message, receivers []string, options ...Option) (status Status, err error) {
	opts := defaultOptions
	if options != nil {
		for _, opt := range options {
			opt.f(&opts)
		}
	}
	if receivers == nil || len(receivers) == 0 {
		// return StatusError, errors.New("receivers is empty")
		return StatusSuccess, nil
	}

	// encode msg
	message, err := protoMessage2Deliver(d.svcName, msg)
	if err != nil {
		return StatusError, err
	}

	// deliver to gateway
	if len(receivers) == 1 {
		// deliver one receiver
		reqDeliver := &postal.ReqDeliver{
			Receiver: receivers[0],
			Msg:      message,
		}

		ctx = context.WithValue(ctx, balancer.ConsistentHashKey, reqDeliver.Receiver)
		res, err := d.postal.Deliver(ctx, reqDeliver)
		if err != nil {
			return StatusError, err
		}
		// TODO res code
		if res.Ok {
			status = StatusSuccess
		} else {
			status = StatusError
		}
		logger.Info("deliver result: ", err, res)
	} else {

		// deliver batch receiver
		ctx = context.WithValue(ctx, balancer.ConsistentHashKey, receivers[0])
		req := &postal.ReqDeliverBatch{Receivers: receivers, Msg: message}
		res, err := d.postal.DeliverBatch(ctx, req)
		if err != nil {
			logger.Error("deliver error: ", err)
			return StatusError, err
		}
		if res.Ok {
			status = StatusSuccess
		} else {
			status = StatusError
		}
		logger.Info("deliver batch result: ", res)
	}
	return
}