TUN-5142: Add asynchronous servecontrolstream for QUIC

ServeControlStream accidentally became non-blocking in the last quic
change causing stream to not be returned until a SIGTERM was received.
This change makes ServeControlStream be non-blocking for QUIC streams.

connection/control.go

@@ -29,7 +29,7 @@ type controlStream struct {
 
 // ControlStreamHandler registers connections with origintunneld and initiates graceful shutdown.
 type ControlStreamHandler interface {
-	ServeControlStream(ctx context.Context, rw io.ReadWriteCloser, connOptions *tunnelpogs.ConnectionOptions) error
+	ServeControlStream(ctx context.Context, rw io.ReadWriteCloser, connOptions *tunnelpogs.ConnectionOptions, shouldWaitForUnregister bool) error
 	IsStopped() bool
 }
 
@@ -61,6 +61,7 @@ func (c *controlStream) ServeControlStream(
 	ctx context.Context,
 	rw io.ReadWriteCloser,
 	connOptions *tunnelpogs.ConnectionOptions,
+	shouldWaitForUnregister bool,
 ) error {
 	rpcClient := c.newRPCClientFunc(ctx, rw, c.observer.log)
 	defer rpcClient.Close()
@@ -70,6 +71,16 @@ func (c *controlStream) ServeControlStream(
 	}
 	c.connectedFuse.Connected()
 
+	if shouldWaitForUnregister {
+		c.waitForUnregister(ctx, rpcClient)
+	} else {
+		go c.waitForUnregister(ctx, rpcClient)
+	}
+
+	return nil
+}
+
+func (c *controlStream) waitForUnregister(ctx context.Context, rpcClient NamedTunnelRPCClient) {
 	// wait for connection termination or start of graceful shutdown
 	select {
 	case <-ctx.Done():
@@ -81,7 +92,6 @@ func (c *controlStream) ServeControlStream(
 	c.observer.sendUnregisteringEvent(c.connIndex)
 	rpcClient.GracefulShutdown(ctx, c.gracePeriod)
 	c.observer.log.Info().Uint8(LogFieldConnIndex, c.connIndex).Msg("Unregistered tunnel connection")
-	return nil
 }
 
 func (c *controlStream) IsStopped() bool {