137 lines
2.9 KiB
Go
137 lines
2.9 KiB
Go
package server
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
pb "gitlab.michelsen.id/phillmichelsen/tessera/pkg/pb/data_service"
|
|
"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
type SocketStreamingServer struct {
|
|
manager *manager.Manager
|
|
}
|
|
|
|
func NewSocketStreamingServer(m *manager.Manager) *SocketStreamingServer {
|
|
return &SocketStreamingServer{manager: m}
|
|
}
|
|
|
|
func (s *SocketStreamingServer) Serve(lis net.Listener) error {
|
|
for {
|
|
conn, err := lis.Accept()
|
|
if err != nil {
|
|
fmt.Printf("accept error: %v\n", err)
|
|
continue
|
|
}
|
|
go s.handleConnection(conn)
|
|
}
|
|
}
|
|
|
|
func (s *SocketStreamingServer) handleConnection(conn net.Conn) {
|
|
defer func() {
|
|
if err := conn.Close(); err != nil {
|
|
fmt.Printf("conn close error: %v\n", err)
|
|
} else {
|
|
fmt.Println("connection closed")
|
|
}
|
|
}()
|
|
|
|
if tc, ok := conn.(*net.TCPConn); ok {
|
|
_ = tc.SetNoDelay(true)
|
|
_ = tc.SetWriteBuffer(512 * 1024)
|
|
_ = tc.SetReadBuffer(512 * 1024)
|
|
_ = tc.SetKeepAlive(true)
|
|
_ = tc.SetKeepAlivePeriod(30 * time.Second)
|
|
}
|
|
|
|
reader := bufio.NewReader(conn)
|
|
|
|
raw, err := reader.ReadString('\n')
|
|
if err != nil {
|
|
fmt.Printf("read stream UUID error: %v\n", err)
|
|
_, _ = fmt.Fprint(conn, "Failed to read stream UUID\n")
|
|
return
|
|
}
|
|
streamUUIDStr := strings.TrimSpace(raw)
|
|
streamUUID, err := uuid.Parse(streamUUIDStr)
|
|
if err != nil {
|
|
_, _ = fmt.Fprint(conn, "Invalid stream UUID\n")
|
|
return
|
|
}
|
|
|
|
outCh, err := s.manager.ConnectStream(streamUUID)
|
|
if err != nil {
|
|
_, _ = fmt.Fprintf(conn, "Failed to connect to stream: %v\n", err)
|
|
return
|
|
}
|
|
defer s.manager.DisconnectStream(streamUUID)
|
|
|
|
writer := bufio.NewWriterSize(conn, 256*1024)
|
|
defer func(w *bufio.Writer) {
|
|
if err := w.Flush(); err != nil {
|
|
fmt.Printf("final flush error: %v\n", err)
|
|
}
|
|
}(writer)
|
|
|
|
const flushEvery = 32
|
|
batch := 0
|
|
|
|
for msg := range outCh {
|
|
m := pb.Message{
|
|
Identifier: &pb.Identifier{Key: msg.Identifier.Key()},
|
|
Payload: msg.Payload,
|
|
Encoding: string(msg.Encoding),
|
|
}
|
|
|
|
size := proto.Size(&m)
|
|
buf := make([]byte, 0, size)
|
|
b, err := proto.MarshalOptions{}.MarshalAppend(buf, &m)
|
|
if err != nil {
|
|
fmt.Printf("proto marshal error: %v\n", err)
|
|
continue
|
|
}
|
|
|
|
var hdr [4]byte
|
|
if len(b) > int(^uint32(0)) {
|
|
fmt.Printf("message too large: %d bytes\n", len(b))
|
|
continue
|
|
}
|
|
binary.BigEndian.PutUint32(hdr[:], uint32(len(b)))
|
|
|
|
if _, err := writer.Write(hdr[:]); err != nil {
|
|
if err == io.EOF {
|
|
return
|
|
}
|
|
fmt.Printf("write len error: %v\n", err)
|
|
return
|
|
}
|
|
if _, err := writer.Write(b); err != nil {
|
|
if err == io.EOF {
|
|
return
|
|
}
|
|
fmt.Printf("write body error: %v\n", err)
|
|
return
|
|
}
|
|
|
|
batch++
|
|
if batch >= flushEvery {
|
|
if err := writer.Flush(); err != nil {
|
|
fmt.Printf("flush error: %v\n", err)
|
|
return
|
|
}
|
|
batch = 0
|
|
}
|
|
}
|
|
|
|
if err := writer.Flush(); err != nil {
|
|
fmt.Printf("final flush error: %v\n", err)
|
|
}
|
|
}
|