From b0e588d82261ea81b3f6be8fa009136e2ba9d31e Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 25 Jan 2024 14:10:56 +0700
Subject: [PATCH] schedule sending when writing initial or handshake crypto
 data

No functional change expected. We already did schedule sending for all
cases but the ClientHello. Scheduling explicitly allows us to remove
this corner case.
---
 connection.go      | 6 ++----
 connection_test.go | 1 +
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/connection.go b/connection.go
index 9c8821d91ff..cb31e143f97 100644
--- a/connection.go
+++ b/connection.go
@@ -514,10 +514,6 @@ func (s *connection) run() error {
 		}
 	}()
 
-	if s.perspective == protocol.PerspectiveClient {
-		s.scheduleSending() // so the ClientHello actually gets sent
-	}
-
 	var sendQueueAvailable <-chan struct{}
 
 runLoop:
@@ -1406,8 +1402,10 @@ func (s *connection) handleHandshakeEvents() error {
 			err = s.dropEncryptionLevel(protocol.Encryption0RTT)
 		case handshake.EventWriteInitialData:
 			_, err = s.initialStream.Write(ev.Data)
+			s.scheduleSending()
 		case handshake.EventWriteHandshakeData:
 			_, err = s.handshakeStream.Write(ev.Data)
+			s.scheduleSending()
 		}
 		if err != nil {
 			return err
diff --git a/connection_test.go b/connection_test.go
index fe318d22171..26679987185 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -2588,6 +2588,7 @@ var _ = Describe("Client Connection", func() {
 			defer GinkgoRecover()
 			cryptoSetup.EXPECT().StartHandshake().MaxTimes(1)
 			cryptoSetup.EXPECT().NextEvent().Return(handshake.Event{Kind: handshake.EventNoEvent})
+			conn.scheduleSending()
 			conn.run()
 		}()
 		Eventually(done).Should(BeClosed())