From 0c561f23bf828e5c449a3f77481e1490a2f6ec49 Mon Sep 17 00:00:00 2001
From: patterniha <71074308+patterniha@users.noreply.github.com>
Date: Tue, 15 Jul 2025 06:11:44 +0330
Subject: [PATCH] Update dispatcher.go

---
 transport/internet/udp/dispatcher.go | 37 +++++++++++++++++++---------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/transport/internet/udp/dispatcher.go b/transport/internet/udp/dispatcher.go
index e6267d2f..1dea6de7 100644
--- a/transport/internet/udp/dispatcher.go
+++ b/transport/internet/udp/dispatcher.go
@@ -24,6 +24,15 @@ type connEntry struct {
 	link   *transport.Link
 	timer  signal.ActivityUpdater
 	cancel context.CancelFunc
+	closed bool
+}
+
+func (c *connEntry) Close() error {
+	c.closed = true
+	c.cancel()
+	common.Interrupt(c.link.Reader)
+	common.Close(c.link.Writer)
+	return nil
 }
 
 type Dispatcher struct {
@@ -45,8 +54,7 @@ func (v *Dispatcher) RemoveRay() {
 	v.Lock()
 	defer v.Unlock()
 	if v.conn != nil {
-		common.Interrupt(v.conn.link.Reader)
-		common.Close(v.conn.link.Writer)
+		v.conn.Close()
 		v.conn = nil
 	}
 }
@@ -56,28 +64,32 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
 	defer v.Unlock()
 
 	if v.conn != nil {
-		return v.conn, nil
+		if v.conn.closed {
+			v.conn = nil
+		} else {
+			return v.conn, nil
+		}
 	}
 
 	errors.LogInfo(ctx, "establishing new connection for ", dest)
 
 	ctx, cancel := context.WithCancel(ctx)
-	removeRay := func() {
-		cancel()
-		v.RemoveRay()
-	}
-	timer := signal.CancelAfterInactivity(ctx, removeRay, time.Minute)
 
 	link, err := v.dispatcher.Dispatch(ctx, dest)
 	if err != nil {
+		cancel()
 		return nil, errors.New("failed to dispatch request to ", dest).Base(err)
 	}
 
 	entry := &connEntry{
 		link:   link,
-		timer:  timer,
-		cancel: removeRay,
+		cancel: cancel,
 	}
+	//entryClose := func() {
+	//	entry.Close()
+	//}
+
+	entry.timer = signal.CancelAfterInactivity(ctx, cancel, time.Minute)
 	v.conn = entry
 	go handleInput(ctx, entry, dest, v.callback, v.callClose)
 	return entry, nil
@@ -96,7 +108,7 @@ func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination,
 	if outputStream != nil {
 		if err := outputStream.WriteMultiBuffer(buf.MultiBuffer{payload}); err != nil {
 			errors.LogInfoInner(ctx, err, "failed to write first UDP payload")
-			conn.cancel()
+			conn.Close()
 			return
 		}
 	}
@@ -104,7 +116,7 @@ func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination,
 
 func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, callback ResponseCallback, callClose func() error) {
 	defer func() {
-		conn.cancel()
+		conn.Close()
 		if callClose != nil {
 			callClose()
 		}
@@ -136,6 +148,7 @@ func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, cal
 				Payload: b,
 				Source:  dest,
 			})
+			b.Release()
 		}
 	}
 }