feat: atomic boolean check for isClosed

agent.go

@@ -122,6 +122,7 @@ type Agent struct {
 
 	// State for closing
 	done         chan struct{}
+	closed       atomic.Bool // Atomic bools are much faster than channels
 	taskLoopDone chan struct{}
 	err          atomicx.Error
 
@@ -170,10 +171,13 @@ func (a *Agent) getAfterRunFn() []func(context.Context) {
 }
 
 func (a *Agent) ok() error {
-	select {
-	case <-a.done:
+	// select {
+	// case <-a.done:
+	// 	return a.getErr()
+	// default:
+	// }
+	if a.closed.Load() {
 		return a.getErr()
-	default:
 	}
 	return nil
 }
@@ -936,6 +940,7 @@ func (a *Agent) Close() error {
 
 	a.removeUfragFromMux()
 
+	a.closed.Store(true)
 	close(a.done)
 	<-a.taskLoopDone
 	return nil

udp_muxed_conn.go

@@ -36,6 +36,7 @@ type udpMuxedConn struct {
 	// Channel holding incoming packets
 	buf        *packetio.Buffer
 	closedChan chan struct{}
+	closed     atomic.Bool
 	closeOnce  sync.Once
 	mu         sync.Mutex
 
@@ -131,17 +132,19 @@ func (c *udpMuxedConn) Close() error {
 	c.closeOnce.Do(func() {
 		err = c.buf.Close()
 		close(c.closedChan)
+		c.closed.Store(true)
 	})
 	return err
 }
 
 func (c *udpMuxedConn) isClosed() bool {
-	select {
-	case <-c.closedChan:
-		return true
-	default:
-		return false
-	}
+	// select {
+	// case <-c.closedChan:
+	// 	return true
+	// default:
+	// 	return false
+	// }
+	return c.closed.Load()
 }
 
 func (c *udpMuxedConn) getAddresses() []string {