astontest/pkg/messenger/service.go

package messenger

import (
	"context"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"

	"git.derfenix.pro/fenix/astontest/api"
)

type DataStorage interface {
	StoreReceived(ctx context.Context, sender, message string) error
	StoreSent(ctx context.Context, sender, message string) error
}

type DataSource interface {
	NewString(ctx context.Context) string
}

func NewMessenger(newNodes <-chan string, port uint16, storage DataStorage, source DataSource, pingInterval time.Duration, log *zap.Logger) *Messenger {
	return &Messenger{
		newNodes:     newNodes,
		storage:      storage,
		source:       source,
		port:         port,
		log:          log,
		pingInterval: pingInterval,
		FailNodesCh:  make(chan string, 1),
	}
}

type Messenger struct {
	newNodes     <-chan string
	storage      DataStorage
	source       DataSource
	port         uint16
	log          *zap.Logger
	pingInterval time.Duration

	FailNodesCh chan string

	api.UnimplementedAstonTestServer
}

func (m *Messenger) Start(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
	defer close(m.FailNodesCh)

	for {
		select {
		case <-ctx.Done():
			return

		case node := <-m.newNodes:
			wg.Add(1)
			go m.newStream(ctx, wg, node)
		}
	}
}

func (m *Messenger) StartServer(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()

	srv := grpc.NewServer(grpc.Creds(insecure.NewCredentials()))
	api.RegisterAstonTestServer(srv, m)

	lis, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", m.port))
	if err != nil {
		m.log.Error("failed to create listener", zap.Error(err))

		return
	}

	wg.Add(1)
	go func() {
		<-ctx.Done()

		go func() {
			time.Sleep(time.Second)
			srv.Stop()
		}()

		srv.GracefulStop()
		wg.Done()
	}()

	m.log.Info("server start")

	if err := srv.Serve(lis); err != nil {
		m.log.Error("serve error", zap.Error(err))
	}
}

func (m *Messenger) Ping(stream api.AstonTest_PingServer) error {
	ticker := time.NewTicker(m.pingInterval)
	defer ticker.Stop()

	peerData, peerFound := peer.FromContext(stream.Context())
	if !peerFound {
		m.log.Warn("not peer found in context")
		peerData = &peer.Peer{
			Addr: &net.IPAddr{},
		}
	}

	peerIP, _, _ := strings.Cut(peerData.Addr.String(), ":")

	for {
		select {
		case <-stream.Context().Done():
			return nil

		case <-ticker.C:
			data := m.source.NewString(stream.Context())

			if err := stream.Send(&api.PingPong{Value: data}); err != nil {
				m.log.Error("failed to send message", zap.Error(err))

				if code := status.Convert(err).Code(); code == codes.Canceled || code == codes.Unavailable {
					return nil
				}

				continue
			}

			if err := m.storage.StoreSent(stream.Context(), peerIP, data); err != nil {
				m.log.Error("failed to save sent data", zap.Error(err), zap.String("data", data), zap.String("remote_addr", peerIP))
			}
		}
	}
}

func (m *Messenger) newStream(ctx context.Context, wg *sync.WaitGroup, address string) {
	defer wg.Done()

	defer func() {
		m.FailNodesCh <- address
	}()

	log := m.log.With(zap.String("address", address))

	cc, err := grpc.Dial(fmt.Sprintf("%s:%d", address, m.port), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Error("failed to dial", zap.Error(err))

		return
	}

	defer func() {
		if err := cc.Close(); err != nil {
			log.Error("close connection", zap.Error(err))
		}
	}()

	client := api.NewAstonTestClient(cc)
	pingClient, err := client.Ping(ctx)
	if err != nil {
		log.Error("create ping stream", zap.Error(err))

		return
	}

	log.Info("stream started")

	msgCh := m.newMessagesCh(pingClient, log)

	for {
		select {
		case <-ctx.Done():
			return

		case msg, ok := <-msgCh:
			if !ok {
				log.Warn("message channel closed")

				return
			}

			if msg.GetValue() == "" {
				continue
			}

			log.Debug("received new message", zap.String("message", msg.GetValue()))

			if err := m.storage.StoreReceived(ctx, address, msg.GetValue()); err != nil {
				log.Error("store message", zap.Error(err), zap.String("message", msg.GetValue()))
			}
		}
	}
}

func (m *Messenger) newMessagesCh(pingClient api.AstonTest_PingClient, log *zap.Logger) chan *api.PingPong {
	msgCh := make(chan *api.PingPong)

	go func() {
		for {
			msg, err := pingClient.Recv()
			if err != nil {
				if code := status.Convert(err).Code(); !(code == codes.Canceled || code == codes.Unavailable) {
					log.Error("message receive failed", zap.Error(err))
				}

				break
			}
			msgCh <- msg
		}
		close(msgCh)
	}()

	return msgCh
}