NOISSUE - Reduce message loss via vsock with acks (#252)

* state check within func

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* debug logs sending

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* debug message sending

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* ack messages

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* handle proto better

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* improve concurrency

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* improve manager handling

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* remove debug lines

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* sync next id

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* reduce locks

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
This commit is contained in:
Sammy Kerata Oina
2024-09-23 19:38:02 +03:00
committed by GitHub
parent df923f9b1f
commit 5d5ae35e2b
7 changed files with 297 additions and 48 deletions
+29 -21
View File
@@ -8,7 +8,6 @@ package main
import (
"encoding/json"
"encoding/pem"
"fmt"
"log"
"net"
"os"
@@ -17,10 +16,15 @@ import (
"github.com/mdlayher/vsock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/internal"
internalvsock "github.com/ultravioletrs/cocos/internal/vsock"
"github.com/ultravioletrs/cocos/manager"
"github.com/ultravioletrs/cocos/manager/qemu"
pkgmanager "github.com/ultravioletrs/cocos/pkg/manager"
"google.golang.org/protobuf/proto"
)
const (
managerVsockPort = manager.ManagerVsockPort
vsockConfigPort = qemu.VsockConfigPort
)
func main() {
@@ -50,10 +54,6 @@ func main() {
log.Fatalf("failed to calculate checksum: %s", err)
}
l, err := vsock.Listen(manager.ManagerVsockPort, nil)
if err != nil {
log.Fatal(err)
}
ac := agent.Computation{
ID: "123",
Datasets: agent.Datasets{agent.Dataset{Hash: [32]byte(dataHash), UserKey: pubPem.Bytes}},
@@ -65,21 +65,30 @@ func main() {
AttestedTls: attestedTLS,
},
}
if err := SendAgentConfig(3, ac); err != nil {
if err := sendAgentConfig(3, ac); err != nil {
log.Fatal(err)
}
listener, err := vsock.Listen(managerVsockPort, nil)
if err != nil {
log.Fatalf("failed to listen on vsock: %s", err)
}
defer listener.Close()
log.Printf("Listening on vsock port %d", managerVsockPort)
for {
conn, err := l.Accept()
conn, err := listener.Accept()
if err != nil {
log.Println(err)
log.Printf("failed to accept connection: %s", err)
continue
}
go handleConnections(conn)
go handleConnection(conn)
}
}
func SendAgentConfig(cid uint32, ac agent.Computation) error {
func sendAgentConfig(cid uint32, ac agent.Computation) error {
conn, err := vsock.Dial(cid, qemu.VsockConfigPort, nil)
if err != nil {
return err
@@ -100,20 +109,19 @@ func SendAgentConfig(cid uint32, ac agent.Computation) error {
return nil
}
func handleConnections(conn net.Conn) {
func handleConnection(conn net.Conn) {
defer conn.Close()
ackReader := internalvsock.NewAckReader(conn)
for {
b := make([]byte, 1024)
n, err := conn.Read(b)
if err != nil {
log.Println(err)
return
}
var message pkgmanager.ClientStreamMessage
if err := proto.Unmarshal(b[:n], &message); err != nil {
log.Println(err)
err := ackReader.ReadProto(&message)
if err != nil {
log.Printf("Error reading message: %v", err)
return
}
fmt.Println(message.String())
log.Printf("Received message: %s", message.String())
}
}