@@ -69,6 +69,11 @@ private sealed class State
69
69
// Resettable completions to be used for multiple calls to send.
70
70
public readonly ResettableCompletionSource < uint > SendResettableCompletionSource = new ResettableCompletionSource < uint > ( ) ;
71
71
72
+ public ShutdownWriteState ShutdownWriteState ;
73
+
74
+ // Set once writes have been shutdown.
75
+ public readonly TaskCompletionSource ShutdownWriteCompletionSource = new TaskCompletionSource ( TaskCreationOptions . RunContinuationsAsynchronously ) ;
76
+
72
77
public ShutdownState ShutdownState ;
73
78
// The value makes sure that we release the handles only once.
74
79
public int ShutdownDone ;
@@ -577,12 +582,26 @@ internal override void AbortWrite(long errorCode)
577
582
return ;
578
583
}
579
584
585
+ bool shouldComplete = false ;
586
+
580
587
lock ( _state )
581
588
{
582
589
if ( _state . SendState < SendState . Aborted )
583
590
{
584
591
_state . SendState = SendState . Aborted ;
585
592
}
593
+
594
+ if ( _state . ShutdownWriteState == ShutdownWriteState . None )
595
+ {
596
+ _state . ShutdownWriteState = ShutdownWriteState . Canceled ;
597
+ shouldComplete = true ;
598
+ }
599
+ }
600
+
601
+ if ( shouldComplete )
602
+ {
603
+ _state . ShutdownWriteCompletionSource . SetException (
604
+ ExceptionDispatchInfo . SetCurrentStackTrace ( new QuicOperationAbortedException ( "Write was aborted." ) ) ) ;
586
605
}
587
606
588
607
StartShutdown ( QUIC_STREAM_SHUTDOWN_FLAGS . ABORT_SEND , errorCode ) ;
@@ -629,6 +648,23 @@ internal override async ValueTask ShutdownCompleted(CancellationToken cancellati
629
648
await _state . ShutdownCompletionSource . Task . ConfigureAwait ( false ) ;
630
649
}
631
650
651
+ internal override ValueTask WaitForWriteCompletionAsync ( CancellationToken cancellationToken = default )
652
+ {
653
+ // TODO: What should happen if this is called for a unidirectional stream and there are no writes?
654
+
655
+ ThrowIfDisposed ( ) ;
656
+
657
+ lock ( _state )
658
+ {
659
+ if ( _state . ShutdownWriteState == ShutdownWriteState . ConnectionClosed )
660
+ {
661
+ throw GetConnectionAbortedException ( _state ) ;
662
+ }
663
+ }
664
+
665
+ return new ValueTask ( _state . ShutdownWriteCompletionSource . Task . WaitAsync ( cancellationToken ) ) ;
666
+ }
667
+
632
668
internal override void Shutdown ( )
633
669
{
634
670
ThrowIfDisposed ( ) ;
@@ -861,6 +897,11 @@ private static uint HandleEvent(State state, ref StreamEvent evt)
861
897
// Peer has stopped receiving data, don't send anymore.
862
898
case QUIC_STREAM_EVENT_TYPE . PEER_RECEIVE_ABORTED :
863
899
return HandleEventPeerRecvAborted ( state , ref evt ) ;
900
+ // Occurs when shutdown is completed for the send side.
901
+ // This only happens for shutdown on sending, not receiving
902
+ // Receive shutdown can only be abortive.
903
+ case QUIC_STREAM_EVENT_TYPE . SEND_SHUTDOWN_COMPLETE :
904
+ return HandleEventSendShutdownComplete ( state , ref evt ) ;
864
905
// Shutdown for both sending and receiving is completed.
865
906
case QUIC_STREAM_EVENT_TYPE . SHUTDOWN_COMPLETE :
866
907
return HandleEventShutdownComplete ( state , ref evt ) ;
@@ -993,23 +1034,37 @@ private static unsafe uint HandleEventRecv(State state, ref StreamEvent evt)
993
1034
994
1035
private static uint HandleEventPeerRecvAborted ( State state , ref StreamEvent evt )
995
1036
{
996
- bool shouldComplete = false ;
1037
+ bool shouldSendComplete = false ;
1038
+ bool shouldShutdownWriteComplete = false ;
997
1039
lock ( state )
998
1040
{
999
1041
if ( state . SendState == SendState . None || state . SendState == SendState . Pending )
1000
1042
{
1001
- shouldComplete = true ;
1043
+ shouldSendComplete = true ;
1044
+ }
1045
+
1046
+ if ( state . ShutdownWriteState == ShutdownWriteState . None )
1047
+ {
1048
+ state . ShutdownWriteState = ShutdownWriteState . Canceled ;
1049
+ shouldShutdownWriteComplete = true ;
1002
1050
}
1051
+
1003
1052
state . SendState = SendState . Aborted ;
1004
1053
state . SendErrorCode = ( long ) evt . Data . PeerReceiveAborted . ErrorCode ;
1005
1054
}
1006
1055
1007
- if ( shouldComplete )
1056
+ if ( shouldSendComplete )
1008
1057
{
1009
1058
state . SendResettableCompletionSource . CompleteException (
1010
1059
ExceptionDispatchInfo . SetCurrentStackTrace ( new QuicStreamAbortedException ( state . SendErrorCode ) ) ) ;
1011
1060
}
1012
1061
1062
+ if ( shouldShutdownWriteComplete )
1063
+ {
1064
+ state . ShutdownWriteCompletionSource . SetException (
1065
+ ExceptionDispatchInfo . SetCurrentStackTrace ( new QuicStreamAbortedException ( state . SendErrorCode ) ) ) ;
1066
+ }
1067
+
1013
1068
return MsQuicStatusCodes . Success ;
1014
1069
}
1015
1070
@@ -1021,6 +1076,38 @@ private static uint HandleEventStartComplete(State state, ref StreamEvent evt)
1021
1076
return MsQuicStatusCodes . Success ;
1022
1077
}
1023
1078
1079
+ private static uint HandleEventSendShutdownComplete ( State state , ref StreamEvent evt )
1080
+ {
1081
+ // Graceful will be false in three situations:
1082
+ // 1. The peer aborted reads and the PEER_RECEIVE_ABORTED event was raised.
1083
+ // ShutdownWriteCompletionSource is already complete with an error.
1084
+ // 2. We aborted writes.
1085
+ // ShutdownWriteCompletionSource is already complete with an error.
1086
+ // 3. The connection was closed.
1087
+ // SHUTDOWN_COMPLETE event will be raised immediately after this event. It will handle completing with an error.
1088
+ //
1089
+ // Only use this event with sends gracefully completed.
1090
+ if ( evt . Data . SendShutdownComplete . Graceful != 0 )
1091
+ {
1092
+ bool shouldComplete = false ;
1093
+ lock ( state )
1094
+ {
1095
+ if ( state . ShutdownWriteState == ShutdownWriteState . None )
1096
+ {
1097
+ state . ShutdownWriteState = ShutdownWriteState . Finished ;
1098
+ shouldComplete = true ;
1099
+ }
1100
+ }
1101
+
1102
+ if ( shouldComplete )
1103
+ {
1104
+ state . ShutdownWriteCompletionSource . SetResult ( ) ;
1105
+ }
1106
+ }
1107
+
1108
+ return MsQuicStatusCodes . Success ;
1109
+ }
1110
+
1024
1111
private static uint HandleEventShutdownComplete ( State state , ref StreamEvent evt )
1025
1112
{
1026
1113
StreamEventDataShutdownComplete shutdownCompleteEvent = evt . Data . ShutdownComplete ;
@@ -1031,6 +1118,7 @@ private static uint HandleEventShutdownComplete(State state, ref StreamEvent evt
1031
1118
}
1032
1119
1033
1120
bool shouldReadComplete = false ;
1121
+ bool shouldShutdownWriteComplete = false ;
1034
1122
bool shouldShutdownComplete = false ;
1035
1123
1036
1124
lock ( state )
@@ -1040,6 +1128,15 @@ private static uint HandleEventShutdownComplete(State state, ref StreamEvent evt
1040
1128
1041
1129
shouldReadComplete = CleanupReadStateAndCheckPending ( state , ReadState . ReadsCompleted ) ;
1042
1130
1131
+ if ( state . ShutdownWriteState == ShutdownWriteState . None )
1132
+ {
1133
+ // TODO: We can get to this point if the stream is unidirectional and there are no writes.
1134
+ // Consider what is the best behavior here with write shutdown and the read side of
1135
+ // unidirecitonal streams in the future.
1136
+ state . ShutdownWriteState = ShutdownWriteState . Finished ;
1137
+ shouldShutdownWriteComplete = true ;
1138
+ }
1139
+
1043
1140
if ( state . ShutdownState == ShutdownState . None )
1044
1141
{
1045
1142
state . ShutdownState = ShutdownState . Finished ;
@@ -1052,6 +1149,11 @@ private static uint HandleEventShutdownComplete(State state, ref StreamEvent evt
1052
1149
state . ReceiveResettableCompletionSource . Complete ( 0 ) ;
1053
1150
}
1054
1151
1152
+ if ( shouldShutdownWriteComplete )
1153
+ {
1154
+ state . ShutdownWriteCompletionSource . SetResult ( ) ;
1155
+ }
1156
+
1055
1157
if ( shouldShutdownComplete )
1056
1158
{
1057
1159
state . ShutdownCompletionSource . SetResult ( ) ;
@@ -1361,6 +1463,7 @@ private static uint HandleEventConnectionClose(State state)
1361
1463
1362
1464
bool shouldCompleteRead = false ;
1363
1465
bool shouldCompleteSend = false ;
1466
+ bool shouldCompleteShutdownWrite = false ;
1364
1467
bool shouldCompleteShutdown = false ;
1365
1468
1366
1469
lock ( state )
@@ -1373,6 +1476,12 @@ private static uint HandleEventConnectionClose(State state)
1373
1476
}
1374
1477
state . SendState = SendState . ConnectionClosed ;
1375
1478
1479
+ if ( state . ShutdownWriteState == ShutdownWriteState . None )
1480
+ {
1481
+ shouldCompleteShutdownWrite = true ;
1482
+ }
1483
+ state . ShutdownWriteState = ShutdownWriteState . ConnectionClosed ;
1484
+
1376
1485
if ( state . ShutdownState == ShutdownState . None )
1377
1486
{
1378
1487
shouldCompleteShutdown = true ;
@@ -1392,6 +1501,12 @@ private static uint HandleEventConnectionClose(State state)
1392
1501
ExceptionDispatchInfo . SetCurrentStackTrace ( GetConnectionAbortedException ( state ) ) ) ;
1393
1502
}
1394
1503
1504
+ if ( shouldCompleteShutdownWrite )
1505
+ {
1506
+ state . ShutdownWriteCompletionSource . SetException (
1507
+ ExceptionDispatchInfo . SetCurrentStackTrace ( GetConnectionAbortedException ( state ) ) ) ;
1508
+ }
1509
+
1395
1510
if ( shouldCompleteShutdown )
1396
1511
{
1397
1512
state . ShutdownCompletionSource . SetException (
@@ -1493,6 +1608,14 @@ private enum ReadState
1493
1608
Closed
1494
1609
}
1495
1610
1611
+ private enum ShutdownWriteState
1612
+ {
1613
+ None = 0 ,
1614
+ Canceled ,
1615
+ Finished ,
1616
+ ConnectionClosed
1617
+ }
1618
+
1496
1619
private enum ShutdownState
1497
1620
{
1498
1621
None = 0 ,
0 commit comments