CCCPaste

0acaf898d532476c5a58b11c2f62d5a71190628134a4a...

package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"os"
	"sync"

	pb "dagorokhov/messenger/proto"
	emptypb "google.golang.org/protobuf/types/known/emptypb"
	timestamppb "google.golang.org/protobuf/types/known/timestamppb"

	grpc "google.golang.org/grpc"
)

type messengerServer struct {
	pb.UnimplementedMessengerServerServer
	mutex                sync.Mutex
	streams              map[uint]grpc.ServerStreamingServer[pb.Message]
	wait_go_send         map[uint]chan bool
	finish_read_messages map[uint]chan struct{}
	i                    uint
}

func (s *messengerServer) SendMessage(ctx context.Context, message *pb.Message) (*pb.Reply, error) {
	fmt.Println("got message", message)
	reply := pb.Reply{SendTime: timestamppb.Now()}
	message.SendTime = reply.SendTime

	s.mutex.Lock()
	defer s.mutex.Unlock()

	for i, stream := range s.streams {
		go func(channel chan bool) {
			if err := stream.Send(message); err != nil {
				fmt.Println("error while sending message:", err)
				channel <- true
				return
			}
			channel <- false
		}(s.wait_go_send[i])
	}

	for i, channel := range s.wait_go_send {
		if <-channel == true {
			close(s.finish_read_messages[i])
			delete(s.streams, i)
			delete(s.finish_read_messages, i)
			delete(s.wait_go_send, i)
		}
	}

	return &reply, nil
}

func (s *messengerServer) ReadMessages(_ *emptypb.Empty, stream grpc.ServerStreamingServer[pb.Message]) error {
	s.mutex.Lock()

	fmt.Println("connected to listen")
	i := s.i
	s.i += 1
	s.wait_go_send[i] = make(chan bool)
	s.finish_read_messages[i] = make(chan struct{})
	finish_read_messages := s.finish_read_messages[i]
	s.streams[i] = stream

	s.mutex.Unlock()

	<-finish_read_messages
	return nil
}

func main() {
	fmt.Println("starting server")
	port := os.Getenv("MESSENGER_SERVER_PORT")
	if port == "" {
		port = "51075"
	}

	lis, err := net.Listen("tcp", fmt.Sprintf("localhost:%v", port))
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	var opts []grpc.ServerOption
	grpcServer := grpc.NewServer(opts...)

	server := messengerServer{}
	server.streams = make(map[uint]grpc.ServerStreamingServer[pb.Message])
	server.wait_go_send = make(map[uint]chan bool)
	server.finish_read_messages = make(map[uint]chan struct{})
	pb.RegisterMessengerServerServer(grpcServer, &server)

	grpcServer.Serve(lis)
}