From da6fac4133851a69d9a3815ffab9cacde52f518f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20=22Pisco=22=20Fernandes?=
 <joaocarlos@cloudflare.com>
Date: Wed, 13 Mar 2024 13:30:38 +0000
Subject: [PATCH] TUN-8297: Improve write timeout logging on safe_stream.go

## Summary:
In order to properly monitor what is happening with the new write timeouts that we introduced
in TUN-8244 we need proper logging. Right now we were logging write timeouts when the safe
stream was being closed which didn't make sense because it was miss leading, so this commit
prevents that by adding a flag that allows us to know whether we are closing the stream or not.
---
 quic/safe_stream.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/quic/safe_stream.go b/quic/safe_stream.go
index db347aac..a1575b46 100644
--- a/quic/safe_stream.go
+++ b/quic/safe_stream.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/quic-go/quic-go"
@@ -19,6 +20,7 @@ type SafeStreamCloser struct {
 	stream       quic.Stream
 	writeTimeout time.Duration
 	log          *zerolog.Logger
+	closing      atomic.Bool
 }
 
 func NewSafeStreamCloser(stream quic.Stream, writeTimeout time.Duration, log *zerolog.Logger) *SafeStreamCloser {
@@ -44,27 +46,35 @@ func (s *SafeStreamCloser) Write(p []byte) (n int, err error) {
 	}
 	nBytes, err := s.stream.Write(p)
 	if err != nil {
-		s.handleTimeout(err)
+		s.handleWriteError(err)
 	}
 
 	return nBytes, err
 }
 
 // Handles the timeout error in case it happened, by canceling the stream write.
-func (s *SafeStreamCloser) handleTimeout(err error) {
+func (s *SafeStreamCloser) handleWriteError(err error) {
+	// If we are closing the stream we just ignore any write error.
+	if s.closing.Load() {
+		return
+	}
 	var netErr net.Error
 	if errors.As(err, &netErr) {
 		if netErr.Timeout() {
-			// We don't need to log if what cause the timeout was `no network activity`.
+			// We don't need to log if what cause the timeout was no network activity.
 			if !errors.Is(netErr, &idleTimeoutError) {
 				s.log.Error().Err(netErr).Msg("Closing quic stream due to timeout while writing")
 			}
+			// We need to explicitly cancel the write so that it frees all buffers.
 			s.stream.CancelWrite(0)
 		}
 	}
 }
 
 func (s *SafeStreamCloser) Close() error {
+	// Set this stream to a closing state.
+	s.closing.Store(true)
+
 	// Make sure a possible writer does not block the lock forever. We need it, so we can close the writer
 	// side of the stream safely.
 	_ = s.stream.SetWriteDeadline(time.Now())