From 1cfdd98e187830a794eeb407f47cf0d20e012ef3 Mon Sep 17 00:00:00 2001
From: eyedeekay <hankhill19580@gmail.com>
Date: Thu, 29 May 2025 18:49:54 -0400
Subject: [PATCH] Work on making DatagramConn implement net.Conn

---
 datagram/packetconn.go | 10 ++++------
 datagram/read.go       | 20 +++++++++++++++-----
 datagram/types.go      |  1 +
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/datagram/packetconn.go b/datagram/packetconn.go
index 8c14d127..a95f13d9 100644
--- a/datagram/packetconn.go
+++ b/datagram/packetconn.go
@@ -69,16 +69,14 @@ func (c *DatagramConn) Close() error {
 
 	c.closed = true
 
-	// Close reader and writer
+	// Close reader and writer - these are owned by this connection
 	if c.reader != nil {
 		c.reader.Close()
 	}
 
-	// Close the session
-	if err := c.session.Close(); err != nil {
-		logger.WithError(err).Error("Failed to close session")
-		return oops.Errorf("failed to close datagram connection: %w", err)
-	}
+	// DO NOT close the session - it's a shared resource that may be used by other connections
+	// The session should be closed by the code that created it, not by individual connections
+	// that use it. This follows the principle that the creator owns the resource.
 
 	logger.Debug("Successfully closed DatagramConn")
 	return nil
diff --git a/datagram/read.go b/datagram/read.go
index d9836b20..51b1c5d2 100644
--- a/datagram/read.go
+++ b/datagram/read.go
@@ -46,14 +46,21 @@ func (r *DatagramReader) Close() error {
 	logger.Debug("Closing DatagramReader")
 
 	r.closed = true
-	// Signal termination to receiveLoop and wait for it to exit
+
+	// Signal termination to receiveLoop
 	close(r.closeChan)
 
-	// Give receiveLoop time to detect the close signal and exit
-	// before closing the channels it might be writing to
-	time.Sleep(10 * time.Millisecond)
+	// Wait for receiveLoop to signal it has exited by closing doneChan
+	// This ensures proper synchronization without arbitrary delays
+	select {
+	case <-r.doneChan:
+		// receiveLoop has confirmed it stopped
+	case <-time.After(5 * time.Second):
+		// Timeout protection - log warning but continue cleanup
+		logger.Warn("Timeout waiting for receive loop to stop")
+	}
 
-	// Now safe to close the receiver channels
+	// Now safe to close the receiver channels since receiveLoop has stopped
 	close(r.recvChan)
 	close(r.errorChan)
 
@@ -66,6 +73,9 @@ func (r *DatagramReader) receiveLoop() {
 	logger := log.WithField("session_id", r.session.ID())
 	logger.Debug("Starting receive loop")
 
+	// Ensure we signal completion when this loop exits
+	r.doneChan = make(chan struct{})
+
 	for {
 		// Check for closure in a non-blocking way first
 		select {
diff --git a/datagram/types.go b/datagram/types.go
index 580f4cc1..14357792 100644
--- a/datagram/types.go
+++ b/datagram/types.go
@@ -23,6 +23,7 @@ type DatagramReader struct {
 	recvChan  chan *Datagram
 	errorChan chan error
 	closeChan chan struct{}
+	doneChan  chan struct{}
 	closed    bool
 	mu        sync.RWMutex
 }