fix: manage resources properly (#17)

* fix(simulation): remove client from map when sess closes

* fix(simulation): ensure returns when sess closes

* refactor(simulation): pass ctx from outside

* fix(mcp): propagate ctx from readLoop

* fix(game): do not spawn yet another goroutine for sending inputs

* feat(mcp): implement Session.Closed()

* fix(game): terminate if sess closed

* refactor(simulation): accept context as a param

* fix(simulation): break on context cancellation or timeout

* refactor(simulation)!: remove unused ctx

Author: utilyre

cmd/server/main.go

@@ -22,7 +22,7 @@ func main() {
 	ctx, cancel := cli.NewSignalContext()
 	defer cancel()
 
-	sim, err := simulation.New(ctx, localAddr)
+	sim, err := simulation.New(localAddr)
 	if err != nil {
 		slog.Error("failed to instantiate simulation", "error", err)
 		return

internal/game/game.go

@@ -58,32 +58,14 @@ func New(ctx context.Context, raddr string) (*Game, error) {
 		nextSnapshot:    snapshot{},
 		snapshotLock:    sync.Mutex{},
 	}
-	go g.sendLoop()
-	go g.receiveLoop()
+	go g.receiveLoop(context.Background())
 	return g, nil
 }
 
-func (g *Game) sendLoop() {
-	ticker := time.NewTicker(time.Second / time.Duration(ebiten.TPS()))
-	defer ticker.Stop()
-	for ; ; <-ticker.C {
-		g.inputBufferLock.Lock()
-		data, err := g.inputBuffer.MarshalBinary()
-		g.inputBufferLock.Unlock()
-		if err != nil {
-			slog.Warn("failed to marshal input buffer", "error", err)
-			continue
-		}
-
-		_ = g.sess.TrySend(data)
-	}
-}
-
-func (g *Game) receiveLoop() {
-	ctx := context.Background()
+func (g *Game) receiveLoop(ctx context.Context) {
 	for {
 		data, err := g.sess.Receive(ctx)
-		if errors.Is(err, mcp.ErrClosed) {
+		if errors.Is(err, mcp.ErrClosed) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 			break
 		}
 		if err != nil {
@@ -164,7 +146,16 @@ func (g *Game) Update() error {
 
 	g.inputBufferLock.Lock()
 	g.inputBuffer.Append(input)
+	data, err := g.inputBuffer.MarshalBinary()
 	g.inputBufferLock.Unlock()
+	if err != nil {
+		slog.Warn("failed to marshal input buffer", "error", err)
+		return nil
+	}
+	if g.sess.Closed() {
+		return ebiten.Termination
+	}
+	_ = g.sess.TrySend(data)
 
 	g.snapshotLock.Lock()
 	if !g.nextSnapshot.t.IsZero() {

internal/mcp/listener.go

@@ -334,7 +334,7 @@ func (ln *Listener) readLoop() {
 			continue
 		}
 
-		err = ln.handleDatagram(remote, datagram)
+		err = ln.handleDatagram(context.Background(), remote, datagram)
 		if err != nil {
 			ln.logger.Warn("failed to handle datagram", "error", err)
 			continue
@@ -347,7 +347,7 @@ func (ln *Listener) readLoop() {
 	}
 }
 
-func (ln *Listener) handleDatagram(remote net.Addr, datagram Datagram) error {
+func (ln *Listener) handleDatagram(ctx context.Context, remote net.Addr, datagram Datagram) error {
 	if datagram.Version != version {
 		return fmt.Errorf("version %d: version is not supported", datagram.Version)
 	}
@@ -380,7 +380,7 @@ func (ln *Listener) handleDatagram(remote net.Addr, datagram Datagram) error {
 		sess := ln.sessions[remote.String()]
 		var err error
 		sess.dieOnce.Do(func() {
-			err = sess.partialUncheckedClose(context.TODO())
+			err = sess.partialUncheckedClose(ctx)
 		})
 		if err != nil {
 			ln.sessionCond.L.Unlock()
@@ -559,6 +559,15 @@ func (sess *Session) Close(ctx context.Context) error {
 	return err
 }
 
+func (sess *Session) Closed() bool {
+	select {
+	case _, open := <-sess.die:
+		return !open
+	default:
+		return false
+	}
+}
+
 func (sess *Session) LocalAddr() net.Addr {
 	return sess.local
 }

internal/simulation/simulation.go

@@ -25,7 +25,7 @@ type Simulation struct {
 	lastStateIndex uint32
 }
 
-func New(ctx context.Context, laddr string) (*Simulation, error) {
+func New(laddr string) (*Simulation, error) {
 	houseImg, err := openImage("./assets/house.png")
 	if err != nil {
 		return nil, err
@@ -45,7 +45,7 @@ func New(ctx context.Context, laddr string) (*Simulation, error) {
 		state:          state.State{},
 		lastStateIndex: 0,
 	}
-	go sim.acceptLoop()
+	go sim.acceptLoop(context.Background())
 	return sim, nil
 }
 
@@ -54,21 +54,21 @@ type clientType struct {
 	inputc chan state.Input
 }
 
-func (c clientType) start() {
+func (c clientType) start(ctx context.Context) {
 	logger := slog.With("remote", c.sess.RemoteAddr())
 
-	// TODO: this is a temporary fix for the busy-loop performance issue
-	ticker := time.NewTicker(time.Second / time.Duration(ebiten.TPS()))
-	defer ticker.Stop()
-	// TODO: plan for killing this infinite loop
-	for ; ; <-ticker.C {
-		data, succeeded := c.sess.TryReceive()
-		if !succeeded {
+	for {
+		data, err := c.sess.Receive(ctx)
+		if errors.Is(err, mcp.ErrClosed) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			break
+		}
+		if err != nil {
+			logger.Warn("failed to receive inputs from session", "error", err)
 			continue
 		}
 
 		var buf jitter.Buffer
-		err := buf.UnmarshalBinary(data)
+		err = buf.UnmarshalBinary(data)
 		if err != nil {
 			logger.Warn("failed to unmarshal inputs", "error", err)
 			continue
@@ -78,6 +78,7 @@ func (c clientType) start() {
 			b := make([]byte, 2+4)
 			binary.BigEndian.PutUint16(b, 0 /* type = input ack */)
 			binary.BigEndian.PutUint32(b[2:], indices[len(indices)-1])
+			// i refuse to spawn a new goroutine just to do this
 			_ = c.sess.TrySend(b)
 		}
 
@@ -90,28 +91,35 @@ func (c clientType) start() {
 	}
 }
 
-func (sim *Simulation) acceptLoop() {
-	ctx := context.Background()
+func (sim *Simulation) acceptLoop(ctx context.Context) {
 	for {
 		sess, err := sim.ln.Accept(ctx)
-		if errors.Is(err, mcp.ErrClosed) {
+		if errors.Is(err, mcp.ErrClosed) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 			break
 		}
 		if err != nil {
 			slog.Warn("failed to accept session", "error", err)
 			continue
 		}
 
-		// TODO: remove sess from sessions whenever closed
-
 		client := clientType{
 			sess:   sess,
 			inputc: make(chan state.Input, 1),
 		}
-		go client.start()
+		raddr := sess.RemoteAddr().String()
+		go func() {
+			client.start(context.Background())
+
+			// should not sess.Close() since the only reason client.start()
+			// returns is because sess has closed.
+
+			sim.clientLock.Lock()
+			delete(sim.clients, raddr)
+			sim.clientLock.Unlock()
+		}()
 
 		sim.clientLock.Lock()
-		sim.clients[sess.RemoteAddr().String()] = client
+		sim.clients[raddr] = client
 		sim.clientLock.Unlock()
 	}
 }