Switch test to UpgradeToRocket

Summary: Upstream diffs allows us to enable Thrift stress test for UpgradeToRocket/Rocket protocols (finally!).

Reviewed By: siutsin

Differential Revision: D83116964

fbshipit-source-id: cb87b29985cbd4c37a3a331e1426130542a34401

Author: echistyakov

thrift/lib/go/thrift/stress/server_test.go

@@ -36,15 +36,12 @@ import (
 )
 
 func TestServerStress(t *testing.T) {
-	t.Run("Header", func(t *testing.T) {
-		runStressTest(t, thrift.TransportIDHeader)
+	t.Run("UpgradeToRocket", func(t *testing.T) {
+		runStressTest(t, thrift.TransportIDUpgradeToRocket)
+	})
+	t.Run("Rocket", func(t *testing.T) {
+		runStressTest(t, thrift.TransportIDRocket)
 	})
-	// t.Run("UpgradeToRocket", func(t *testing.T) {
-	// 	runStressTest(t, thrift.TransportIDUpgradeToRocket)
-	// })
-	// t.Run("Rocket", func(t *testing.T) {
-	// 	runStressTest(t, thrift.TransportIDRocket)
-	// })
 }
 
 func runStressTest(t *testing.T, serverTransport thrift.TransportID) {
@@ -159,6 +156,9 @@ func runStressTest(t *testing.T, serverTransport thrift.TransportID) {
 	err = serverEG.Wait()
 	require.ErrorIs(t, err, context.Canceled)
 
+	// A tiny sleep to allow for some lingering goroutines to finish up.
+	time.Sleep(10 * time.Millisecond)
+
 	// Go routine check (after server shutdown)
 	// We shouldn't exceed 10 Go-routines, if we do - something didn't get cleaned up properly.
 	goroutinesAfterServerStop := runtime.NumGoroutine()