Gracefully Closing Go TCP/UDP Proxies: Coordinating Read and Write Loops with Context

While implementing a TCP/UDP proxy for Easegress, the author separated socket I/O into two goroutines— read-loop and write-loop —and used conn.Close() when io.EOF was read. This caused the write loop to panic with use of closed network connection, leading to noisy logs and swallowed errors.

To solve the problem, the author adopted a pattern that couples a context.Context with the socket’s read deadline. When the context is cancelled, a helper goroutine sets the read deadline to time.Now(), forcing the blocked read to return a timeout error. Both loops then detect the cancelled context or a shared closed atomic flag and exit gracefully.

type ReadDeadliner interface {
    SetReadDeadline(t time.Time) error
}

func SetReadDeadlineOnCancel(ctx context.Context, d ReadDeadliner) {
    go func() {
        <-ctx.Done()
        d.SetReadDeadline(time.Now())
    }()
}

func handleConnection(ctx context.Context, conn net.Conn) {
    ctx, cancel := context.WithCancel(ctx)
    defer cancel()
    SetReadDeadlineOnCancel(ctx, conn)
    // rest of code goes here
    for {
        n, err := conn.Read(buf) // blocks waiting for a message
        if err != nil {
            // handle error (including timeout when context cancelled)
        }
        // ... processing ...
        conn.Write(...)
    }
}

When the read loop encounters io.EOF, it sets the shared closed atomic variable and calls conn.SetReadDeadline(time.Now()). The write loop, upon receiving a timeout error, checks the same closed flag and exits without further writes.

This approach avoids the need for explicit locks or channel synchronization between the two goroutines. However, a race can still occur if the write loop experiences an error between the moment the closed flag is set and the deadline is applied. The author notes that converting the closed flag from a bool to a chan struct{} does not solve the issue.

The technique mirrors how Google’s mtail project handles read interruption on context cancellation (see the commit “Simplify pipestream by correctly interrupting a read on context cancel”).

In summary, the recommended workflow is:

Detect io.EOF in the read loop, set a shared closed flag, and immediately invoke conn.SetReadDeadline(time.Now()) to unblock any pending reads.

In the write loop, treat timeout errors as a signal that the connection is being closed; check the closed flag and exit the loop.

Optionally, wrap the deadline‑setting logic in a context‑cancellation helper as shown above to keep the code tidy.

Remaining challenges include handling the narrow window where the write loop may still see a non‑timeout error after the flag is set, and deciding whether additional synchronization is truly necessary.