Repository: quic-go/quic-go
Pull request #5234 (2025-06-21 07:21:57)

handshake: store key update interval in an atomic
We recently changed the way the key update interval is set in tests to use an environment variable. This resolved a race condition that existed in the earlier logic, however, parsing of the environment variable now shows up in benchmark tests.

Using an atomic variable should have a negligible performance impact.

This is similar to what @egonelbre proposed in #5147.


Patch:
diff --git a/internal/handshake/updatable_aead.go b/internal/handshake/updatable_aead.go
index 0aca46fe..415f85de 100644
--- a/internal/handshake/updatable_aead.go
+++ b/internal/handshake/updatable_aead.go
@@ -6,9 +6,7 @@ import (
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
-	"os"
-	"strconv"
-	"testing"
+	"sync/atomic"
 	"time"
 
 	"github.com/quic-go/quic-go/internal/protocol"
@@ -17,14 +15,15 @@ import (
 	"github.com/quic-go/quic-go/logging"
 )
 
-func keyUpdateInterval() uint64 {
-	// Reparsing the environment variable is not very performant, but it's only done in tests.
-	if testing.Testing() {
-		if v, err := strconv.ParseUint(os.Getenv("QUIC_GO_TEST_KEY_UPDATE_INTERVAL"), 10, 64); err == nil {
-			return v
-		}
-	}
-	return protocol.KeyUpdateInterval
+var keyUpdateInterval atomic.Uint64
+
+func init() {
+	keyUpdateInterval.Store(protocol.KeyUpdateInterval)
+}
+
+func SetKeyUpdateInterval(v uint64) (reset func()) {
+	old := keyUpdateInterval.Swap(v)
+	return func() { keyUpdateInterval.Store(old) }
 }
 
 // FirstKeyUpdateInterval is the maximum number of packets we send or receive before initiating the first key update.
@@ -302,11 +301,11 @@ func (a *updatableAEAD) shouldInitiateKeyUpdate() bool {
 			return true
 		}
 	}
-	if a.numRcvdWithCurrentKey >= keyUpdateInterval() {
+	if a.numRcvdWithCurrentKey >= keyUpdateInterval.Load() {
 		a.logger.Debugf("Received %d packets with current key phase. Initiating key update to the next key phase: %d", a.numRcvdWithCurrentKey, a.keyPhase+1)
 		return true
 	}
-	if a.numSentWithCurrentKey >= keyUpdateInterval() {
+	if a.numSentWithCurrentKey >= keyUpdateInterval.Load() {
 		a.logger.Debugf("Sent %d packets with current key phase. Initiating key update to the next key phase: %d", a.numSentWithCurrentKey, a.keyPhase+1)
 		return true
 	}


Test patch:
diff --git a/integrationtests/self/key_update_test.go b/integrationtests/self/key_update_test.go
index d70c1a8e..32a2d377 100644
--- a/integrationtests/self/key_update_test.go
+++ b/integrationtests/self/key_update_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/internal/handshake"
 	"github.com/quic-go/quic-go/internal/protocol"
 	"github.com/quic-go/quic-go/logging"
 
@@ -14,7 +15,8 @@ import (
 )
 
 func TestKeyUpdates(t *testing.T) {
-	t.Setenv("QUIC_GO_TEST_KEY_UPDATE_INTERVAL", "1") // update keys as frequently as possible
+	reset := handshake.SetKeyUpdateInterval(1) // update keys as frequently as possible
+	t.Cleanup(reset)
 
 	var sentHeaders []*logging.ShortHeader
 	var receivedHeaders []*logging.ShortHeader
diff --git a/internal/handshake/updatable_aead_test.go b/internal/handshake/updatable_aead_test.go
index 36e0638e..bfa9e331 100644
--- a/internal/handshake/updatable_aead_test.go
+++ b/internal/handshake/updatable_aead_test.go
@@ -330,7 +330,8 @@ func TestRejectFrequentKeyUpdates(t *testing.T) {
 }
 
 func setKeyUpdateIntervals(t *testing.T, firstKeyUpdateInterval, keyUpdateInterval uint64) {
-	t.Setenv("QUIC_GO_TEST_KEY_UPDATE_INTERVAL", fmt.Sprintf("%d", keyUpdateInterval))
+	reset := SetKeyUpdateInterval(keyUpdateInterval)
+	t.Cleanup(reset)
 
 	origFirstKeyUpdateInterval := FirstKeyUpdateInterval
 	FirstKeyUpdateInterval = firstKeyUpdateInterval