Work on making DatagramConn implement net.Conn

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
+	// DO NOT close the session - it's a shared resource that may be used by other connections
-	if err := c.session.Close(); err != nil {
+	// The session should be closed by the code that created it, not by individual connections
-		logger.WithError(err).Error("Failed to close session")
+	// that use it. This follows the principle that the creator owns the resource.
-		return oops.Errorf("failed to close datagram connection: %w", err)
-	}
 
 	logger.Debug("Successfully closed DatagramConn")
 	return nil

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
+	// Wait for receiveLoop to signal it has exited by closing doneChan
-	// before closing the channels it might be writing to
+	// This ensures proper synchronization without arbitrary delays
-	time.Sleep(10 * time.Millisecond)
+	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 {

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
 }