sonet/pkg/protocol/deliver/deliver.go

package deliver

import (
	"context"
	"errors"
	"google.golang.org/protobuf/proto"
	"reflect"
	"sonet/api/gen/postal"
	"sonet/pkg/utils/logger"
	"time"
)

type Status int16

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

func Proto2DeliverMessage(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
}

// Deliver n包,通知消息投递
type Deliver struct {
	serviceName  string
	postalPicker *PostalPicker
}

func NewDeliver(svcName string, postalPicker *PostalPicker) *Deliver {
	return &Deliver{
		serviceName:  svcName,
		postalPicker: postalPicker,
	}
}

func (d *Deliver) Deliver(ctx context.Context, msg proto.Message, receiver string, options ...Option) (Status, error) {
	client, err := d.postalPicker.Pick(receiver)
	if err != nil {
		return StatusError, err
	}

	message, err := Proto2DeliverMessage(d.serviceName, msg)
	if err != nil {
		logger.Error("encode proto message error: ", err)
		return StatusError, err
	}
	req := &postal.ReqDeliver{Receiver: receiver, Msg: message}
	_, err = client.Deliver(ctx, req)
	if err != nil {
		return StatusError, err
	}
	return StatusSuccess, nil
}

func (d *Deliver) DeliverBatch(ctx context.Context, msg proto.Message, receivers []string, options ...Option) (status Status, err error) {
	// encode msg
	message, err := Proto2DeliverMessage(d.serviceName, msg)
	if err != nil {
		logger.Error("encode proto message error: ", err)
		return StatusError, err
	}

	// 将消息负载均衡pick好后分批发送
	nodeReceivers := make(map[postal.PostalClient][]string, 5)
	for _, receiver := range receivers {
		client, err := d.postalPicker.Pick(receiver)
		if err != nil {
			return StatusError, err
		}
		nodeReceivers[client] = append(nodeReceivers[client], receiver)
	}

	var errs []error
	for node, receivers := range nodeReceivers {
		_, err := node.DeliverBatch(ctx, &postal.ReqDeliverBatch{Msg: message, Receivers: receivers})
		if err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) > 0 {
		err = errors.Join(errs...)
		return StatusError, err
	}
	return StatusSuccess, nil
}