Implement flow control in Objective-C

src/objective-c/GRPCClient/GRPCCall.h

@@ -183,6 +183,8 @@ extern NSString *const kGRPCTrailersKey;
 - (void)didCloseWithTrailingMetadata:(nullable NSDictionary *)trailingMetadata
                                error:(nullable NSError *)error;
 
+- (void)didWriteData;
+
 @end
 
 /**
@@ -263,6 +265,8 @@ extern NSString *const kGRPCTrailersKey;
  */
 - (void)finish;
 
+- (void)receiveNextMessage;
+
 /**
  * Get a copy of the original call options.
  */

src/objective-c/GRPCClient/GRPCCall.m

@@ -63,6 +63,15 @@ const char *kCFStreamVarName = "grpc_cfstream";
               requestsWriter:(GRXWriter *)requestsWriter
                  callOptions:(GRPCCallOptions *)callOptions;
 
+- (instancetype)initWithHost:(NSString *)host
+                        path:(NSString *)path
+                  callSafety:(GRPCCallSafety)safety
+              requestsWriter:(GRXWriter *)requestsWriter
+                 callOptions:(GRPCCallOptions *)callOptions
+                   writeDone:(void (^)(void))writeDone;
+
+- (void)receiveNextMessage;
+
 @end
 
 @implementation GRPCRequestOptions
@@ -190,7 +199,14 @@ const char *kCFStreamVarName = "grpc_cfstream";
                                       path:_requestOptions.path
                                 callSafety:_requestOptions.safety
                             requestsWriter:_pipe
-                               callOptions:_callOptions];
+                               callOptions:_callOptions
+                                 writeDone:^{
+                                   @synchronized(self) {
+                                     if (self->_handler) {
+                                       [self issueDidWriteData];
+                                     }
+                                   }
+                                 }];
     if (_callOptions.initialMetadata) {
       [_call.requestHeaders addEntriesFromDictionary:_callOptions.initialMetadata];
     }
@@ -363,6 +379,28 @@ const char *kCFStreamVarName = "grpc_cfstream";
   }
 }
 
+- (void)issueDidWriteData {
+  @synchronized(self) {
+    if (_callOptions.enableFlowControl && [_handler respondsToSelector:@selector(didWriteData)]) {
+      dispatch_async(_dispatchQueue, ^{
+        id<GRPCResponseHandler> copiedHandler = nil;
+        @synchronized(self) {
+          copiedHandler = self->_handler;
+        };
+        [copiedHandler didWriteData];
+      });
+    }
+  }
+}
+
+- (void)receiveNextMessage {
+  GRPCCall *copiedCall = nil;
+  @synchronized(self) {
+    copiedCall = _call;
+  }
+  [copiedCall receiveNextMessage];
+}
+
 @end
 
 // The following methods of a C gRPC call object aren't reentrant, and thus
@@ -426,6 +464,15 @@ const char *kCFStreamVarName = "grpc_cfstream";
 
   // The OAuth2 token fetched from a token provider.
   NSString *_fetchedOauth2AccessToken;
+
+  // The callback to be called when a write message op is done.
+  void (^_writeDone)(void);
+
+  // Indicate a read request to core is pending.
+  BOOL _pendingCoreRead;
+
+  // Indicate a read message request from user.
+  BOOL _pendingReceiveNextMessage;
 }
 
 @synthesize state = _state;
@@ -482,12 +529,26 @@ const char *kCFStreamVarName = "grpc_cfstream";
 - (instancetype)initWithHost:(NSString *)host
                         path:(NSString *)path
                   callSafety:(GRPCCallSafety)safety
-              requestsWriter:(GRXWriter *)requestWriter
+              requestsWriter:(GRXWriter *)requestsWriter
                  callOptions:(GRPCCallOptions *)callOptions {
+  return [self initWithHost:host
+                       path:path
+                 callSafety:safety
+             requestsWriter:requestsWriter
+                callOptions:callOptions
+                  writeDone:nil];
+}
+
+- (instancetype)initWithHost:(NSString *)host
+                        path:(NSString *)path
+                  callSafety:(GRPCCallSafety)safety
+              requestsWriter:(GRXWriter *)requestsWriter
+                 callOptions:(GRPCCallOptions *)callOptions
+                   writeDone:(void (^)(void))writeDone {
   // Purposely using pointer rather than length (host.length == 0) for backwards compatibility.
   NSAssert(host != nil && path != nil, @"Neither host nor path can be nil.");
   NSAssert(safety <= GRPCCallSafetyCacheableRequest, @"Invalid call safety value.");
-  NSAssert(requestWriter.state == GRXWriterStateNotStarted,
+  NSAssert(requestsWriter.state == GRXWriterStateNotStarted,
            @"The requests writer can't be already started.");
   if (!host || !path) {
     return nil;
@@ -495,7 +556,7 @@ const char *kCFStreamVarName = "grpc_cfstream";
   if (safety > GRPCCallSafetyCacheableRequest) {
     return nil;
   }
-  if (requestWriter.state != GRXWriterStateNotStarted) {
+  if (requestsWriter.state != GRXWriterStateNotStarted) {
     return nil;
   }
 
@@ -508,16 +569,20 @@ const char *kCFStreamVarName = "grpc_cfstream";
     // Serial queue to invoke the non-reentrant methods of the grpc_call object.
     _callQueue = dispatch_queue_create("io.grpc.call", DISPATCH_QUEUE_SERIAL);
 
-    _requestWriter = requestWriter;
-
+    _requestWriter = requestsWriter;
     _requestHeaders = [[GRPCRequestHeaders alloc] initWithCall:self];
+    _writeDone = writeDone;
 
-    if ([requestWriter isKindOfClass:[GRXImmediateSingleWriter class]]) {
+    if ([requestsWriter isKindOfClass:[GRXImmediateSingleWriter class]]) {
       _unaryCall = YES;
       _unaryOpBatch = [NSMutableArray arrayWithCapacity:kMaxClientBatch];
     }
 
     _responseQueue = dispatch_get_main_queue();
+
+    // do not start a read until initial metadata is received
+    _pendingReceiveNextMessage = NO;
+    _pendingCoreRead = YES;
   }
   return self;
 }
@@ -589,25 +654,28 @@ const char *kCFStreamVarName = "grpc_cfstream";
 // If the call is currently paused, this is a noop. Restarting the call will invoke this
 // method.
 // TODO(jcanizales): Rename to readResponseIfNotPaused.
-- (void)startNextRead {
+- (void)maybeStartNextRead {
   @synchronized(self) {
     if (_state != GRXWriterStateStarted) {
       return;
     }
+    if (_callOptions.enableFlowControl && (_pendingCoreRead || !_pendingReceiveNextMessage)) {
+      return;
+    }
   }
 
   dispatch_async(_callQueue, ^{
     __weak GRPCCall *weakSelf = self;
     [self startReadWithHandler:^(grpc_byte_buffer *message) {
-      if (message == NULL) {
-        // No more messages from the server
-        return;
-      }
       __strong GRPCCall *strongSelf = weakSelf;
       if (strongSelf == nil) {
         grpc_byte_buffer_destroy(message);
         return;
       }
+      if (message == NULL) {
+        // No more messages from the server
+        return;
+      }
       NSData *data = [NSData grpc_dataWithByteBuffer:message];
       grpc_byte_buffer_destroy(message);
       if (!data) {
@@ -616,6 +684,7 @@ const char *kCFStreamVarName = "grpc_cfstream";
         // that's on the hands of any server to have. Instead we finish and ask
         // the server to cancel.
         @synchronized(strongSelf) {
+          strongSelf->_pendingReceiveNextMessage = NO;
           [strongSelf
               finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
                                                   code:GRPCErrorCodeResourceExhausted
@@ -631,7 +700,8 @@ const char *kCFStreamVarName = "grpc_cfstream";
         @synchronized(strongSelf) {
           [strongSelf->_responseWriteable enqueueValue:data
                                      completionHandler:^{
-                                       [strongSelf startNextRead];
+                                       strongSelf->_pendingCoreRead = NO;
+                                       [strongSelf maybeStartNextRead];
                                      }];
         }
       }
@@ -682,6 +752,17 @@ const char *kCFStreamVarName = "grpc_cfstream";
   });
 }
 
+- (void)receiveNextMessage {
+  @synchronized(self) {
+    // Duplicate invocation of this method. Return
+    if (_pendingReceiveNextMessage) {
+      return;
+    }
+    _pendingReceiveNextMessage = YES;
+    [self maybeStartNextRead];
+  }
+}
+
 #pragma mark GRXWriteable implementation
 
 // Only called from the call queue. The error handler will be called from the
@@ -695,6 +776,9 @@ const char *kCFStreamVarName = "grpc_cfstream";
     GRPCCall *strongSelf = weakSelf;
     if (strongSelf) {
       strongSelf->_requestWriter.state = GRXWriterStateStarted;
+      if (strongSelf->_writeDone) {
+        strongSelf->_writeDone();
+      }
     }
   };
 
@@ -774,8 +858,11 @@ const char *kCFStreamVarName = "grpc_cfstream";
     // Response headers received.
     __strong GRPCCall *strongSelf = weakSelf;
     if (strongSelf) {
-      strongSelf.responseHeaders = headers;
-      [strongSelf startNextRead];
+      @synchronized(self) {
+        strongSelf.responseHeaders = headers;
+        strongSelf->_pendingCoreRead = NO;
+        [strongSelf maybeStartNextRead];
+      }
     }
   }
       completionHandler:^(NSError *error, NSDictionary *trailers) {
@@ -929,7 +1016,7 @@ const char *kCFStreamVarName = "grpc_cfstream";
       case GRXWriterStateStarted:
         if (_state == GRXWriterStatePaused) {
           _state = newState;
-          [self startNextRead];
+          [self maybeStartNextRead];
         }
         return;
       case GRXWriterStateNotStarted:

src/objective-c/GRPCClient/GRPCCallOptions.h

@@ -90,6 +90,14 @@ typedef NS_ENUM(NSUInteger, GRPCTransportType) {
  */
 @property(readonly) NSTimeInterval timeout;
 
+/**
+ * Enable flow control of a gRPC call. The option is default to NO. If set to YES, writeData: method
+ * should only be called at most once before a didWriteData callback is issued, and
+ * receiveNextMessage: must be called each time before gRPC call issues a didReceiveMessage
+ * callback.
+ */
+@property(readonly) BOOL enableFlowControl;
+
 // OAuth2 parameters. Users of gRPC may specify one of the following two parameters.
 
 /**
@@ -232,6 +240,14 @@ typedef NS_ENUM(NSUInteger, GRPCTransportType) {
  */
 @property(readwrite) NSTimeInterval timeout;
 
+/**
+ * Enable flow control of a gRPC call. The option is default to NO. If set to YES, writeData: method
+ * should only be called at most once before a didWriteData callback is issued, and
+ * receiveNextMessage: must be called each time before gRPC call issues a didReceiveMessage
+ * callback.
+ */
+@property(readwrite) BOOL enableFlowControl;
+
 // OAuth2 parameters. Users of gRPC may specify one of the following two parameters.
 
 /**

src/objective-c/GRPCClient/GRPCCallOptions.m

@@ -22,6 +22,7 @@
 // The default values for the call options.
 static NSString *const kDefaultServerAuthority = nil;
 static const NSTimeInterval kDefaultTimeout = 0;
+static const BOOL kDefaultEnableFlowControl = NO;
 static NSDictionary *const kDefaultInitialMetadata = nil;
 static NSString *const kDefaultUserAgentPrefix = nil;
 static const NSUInteger kDefaultResponseSizeLimit = 0;
@@ -59,6 +60,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
  @protected
   NSString *_serverAuthority;
   NSTimeInterval _timeout;
+  BOOL _enableFlowControl;
   NSString *_oauth2AccessToken;
   id<GRPCAuthorizationProtocol> _authTokenProvider;
   NSDictionary *_initialMetadata;
@@ -84,6 +86,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 
 @synthesize serverAuthority = _serverAuthority;
 @synthesize timeout = _timeout;
+@synthesize enableFlowControl = _enableFlowControl;
 @synthesize oauth2AccessToken = _oauth2AccessToken;
 @synthesize authTokenProvider = _authTokenProvider;
 @synthesize initialMetadata = _initialMetadata;
@@ -109,6 +112,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 - (instancetype)init {
   return [self initWithServerAuthority:kDefaultServerAuthority
                                timeout:kDefaultTimeout
+                     enableFlowControl:kDefaultEnableFlowControl
                      oauth2AccessToken:kDefaultOauth2AccessToken
                      authTokenProvider:kDefaultAuthTokenProvider
                        initialMetadata:kDefaultInitialMetadata
@@ -134,6 +138,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 
 - (instancetype)initWithServerAuthority:(NSString *)serverAuthority
                                 timeout:(NSTimeInterval)timeout
+                      enableFlowControl:(BOOL)enableFlowControl
                       oauth2AccessToken:(NSString *)oauth2AccessToken
                       authTokenProvider:(id<GRPCAuthorizationProtocol>)authTokenProvider
                         initialMetadata:(NSDictionary *)initialMetadata
@@ -158,6 +163,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   if ((self = [super init])) {
     _serverAuthority = [serverAuthority copy];
     _timeout = timeout < 0 ? 0 : timeout;
+    _enableFlowControl = enableFlowControl;
     _oauth2AccessToken = [oauth2AccessToken copy];
     _authTokenProvider = authTokenProvider;
     _initialMetadata =
@@ -193,6 +199,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   GRPCCallOptions *newOptions =
       [[GRPCCallOptions allocWithZone:zone] initWithServerAuthority:_serverAuthority
                                                             timeout:_timeout
+                                                  enableFlowControl:_enableFlowControl
                                                   oauth2AccessToken:_oauth2AccessToken
                                                   authTokenProvider:_authTokenProvider
                                                     initialMetadata:_initialMetadata
@@ -221,6 +228,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   GRPCMutableCallOptions *newOptions = [[GRPCMutableCallOptions allocWithZone:zone]
       initWithServerAuthority:[_serverAuthority copy]
                       timeout:_timeout
+            enableFlowControl:_enableFlowControl
             oauth2AccessToken:[_oauth2AccessToken copy]
             authTokenProvider:_authTokenProvider
               initialMetadata:[[NSDictionary alloc] initWithDictionary:_initialMetadata
@@ -301,6 +309,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 
 @dynamic serverAuthority;
 @dynamic timeout;
+@dynamic enableFlowControl;
 @dynamic oauth2AccessToken;
 @dynamic authTokenProvider;
 @dynamic initialMetadata;
@@ -326,6 +335,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 - (instancetype)init {
   return [self initWithServerAuthority:kDefaultServerAuthority
                                timeout:kDefaultTimeout
+                     enableFlowControl:kDefaultEnableFlowControl
                      oauth2AccessToken:kDefaultOauth2AccessToken
                      authTokenProvider:kDefaultAuthTokenProvider
                        initialMetadata:kDefaultInitialMetadata
@@ -353,6 +363,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   GRPCCallOptions *newOptions =
       [[GRPCCallOptions allocWithZone:zone] initWithServerAuthority:_serverAuthority
                                                             timeout:_timeout
+                                                  enableFlowControl:_enableFlowControl
                                                   oauth2AccessToken:_oauth2AccessToken
                                                   authTokenProvider:_authTokenProvider
                                                     initialMetadata:_initialMetadata
@@ -381,6 +392,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   GRPCMutableCallOptions *newOptions = [[GRPCMutableCallOptions allocWithZone:zone]
       initWithServerAuthority:_serverAuthority
                       timeout:_timeout
+            enableFlowControl:_enableFlowControl
             oauth2AccessToken:_oauth2AccessToken
             authTokenProvider:_authTokenProvider
               initialMetadata:_initialMetadata
@@ -417,6 +429,10 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   }
 }
 
+- (void)setEnableFlowControl:(BOOL)enableFlowControl {
+  _enableFlowControl = enableFlowControl;
+}
+
 - (void)setOauth2AccessToken:(NSString *)oauth2AccessToken {
   _oauth2AccessToken = [oauth2AccessToken copy];
 }

src/objective-c/ProtoRPC/ProtoRPC.h

@@ -57,6 +57,12 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)didCloseWithTrailingMetadata:(nullable NSDictionary *)trailingMetadata
                                error:(nullable NSError *)error;
 
+/**
+ * Issued when a write action is completed. To get a correct flow control behavior, the user of a
+ * call should not make more than one writeMessage: call before receiving this callback.
+ */
+- (void)didWriteMessage;
+
 @end
 
 /** A unary-request RPC call with Protobuf. */
@@ -130,6 +136,8 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)finish;
 
+- (void)receiveNextMessage;
+
 @end
 
 NS_ASSUME_NONNULL_END

src/objective-c/ProtoRPC/ProtoRPC.m

@@ -197,6 +197,14 @@ static NSError *ErrorForBadProto(id proto, Class expectedClass, NSError *parsing
   [copiedCall finish];
 }
 
+- (void)receiveNextMessage {
+  GRPCCall2 *copiedCall;
+  @synchronized(self) {
+    copiedCall = _call;
+  }
+  [copiedCall receiveNextMessage];
+}
+
 - (void)didReceiveInitialMetadata:(NSDictionary *)initialMetadata {
   @synchronized(self) {
     if (initialMetadata != nil &&
@@ -260,6 +268,20 @@ static NSError *ErrorForBadProto(id proto, Class expectedClass, NSError *parsing
   }
 }
 
+- (void)didWriteData {
+  @synchronized(self) {
+    if ([_handler respondsToSelector:@selector(didWriteMessage)]) {
+      dispatch_async(_dispatchQueue, ^{
+        id<GRPCProtoResponseHandler> copiedHandler = nil;
+        @synchronized(self) {
+          copiedHandler = self->_handler;
+        }
+        [copiedHandler didWriteMessage];
+      });
+    }
+  }
+}
+
 - (dispatch_queue_t)dispatchQueue {
   return _dispatchQueue;
 }

src/objective-c/tests/APIv2Tests/APIv2Tests.m

@@ -44,6 +44,7 @@ static GRPCProtoMethod *kFullDuplexCallMethod;
 static const int kSimpleDataLength = 100;
 
 static const NSTimeInterval kTestTimeout = 16;
+static const NSTimeInterval kInvertedTimeout = 2;
 
 // Reveal the _class ivar for testing access
 @interface GRPCCall2 () {
@@ -56,6 +57,11 @@ static const NSTimeInterval kTestTimeout = 16;
 // Convenience class to use blocks as callbacks
 @interface ClientTestsBlockCallbacks : NSObject<GRPCResponseHandler>
 
+- (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
+                                messageCallback:(void (^)(id))messageCallback
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback
+                              writeDataCallback:(void (^)(void))writeDataCallback;
+
 - (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
                                 messageCallback:(void (^)(id))messageCallback
                                   closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback;
@@ -66,21 +72,33 @@ static const NSTimeInterval kTestTimeout = 16;
   void (^_initialMetadataCallback)(NSDictionary *);
   void (^_messageCallback)(id);
   void (^_closeCallback)(NSDictionary *, NSError *);
+  void (^_writeDataCallback)(void);
   dispatch_queue_t _dispatchQueue;
 }
 
 - (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
                                 messageCallback:(void (^)(id))messageCallback
-                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback {
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback
+                              writeDataCallback:(void (^)(void))writeDataCallback {
   if ((self = [super init])) {
     _initialMetadataCallback = initialMetadataCallback;
     _messageCallback = messageCallback;
     _closeCallback = closeCallback;
+    _writeDataCallback = writeDataCallback;
     _dispatchQueue = dispatch_queue_create(nil, DISPATCH_QUEUE_SERIAL);
   }
   return self;
 }
 
+- (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
+                                messageCallback:(void (^)(id))messageCallback
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback {
+  return [self initWithInitialMetadataCallback:initialMetadataCallback
+                               messageCallback:messageCallback
+                                 closeCallback:closeCallback
+                             writeDataCallback:nil];
+}
+
 - (void)didReceiveInitialMetadata:(NSDictionary *)initialMetadata {
   if (self->_initialMetadataCallback) {
     self->_initialMetadataCallback(initialMetadata);
@@ -99,6 +117,12 @@ static const NSTimeInterval kTestTimeout = 16;
   }
 }
 
+- (void)didWriteData {
+  if (self->_writeDataCallback) {
+    self->_writeDataCallback();
+  }
+}
+
 - (dispatch_queue_t)dispatchQueue {
   return _dispatchQueue;
 }
@@ -475,4 +499,219 @@ static const NSTimeInterval kTestTimeout = 16;
   [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
 }
 
+- (void)testWriteFlowControl {
+  __weak XCTestExpectation *expectWriteData =
+      [self expectationWithDescription:@"Reported write data"];
+
+  RMTStreamingOutputCallRequest *request = [RMTStreamingOutputCallRequest message];
+  RMTResponseParameters *parameters = [RMTResponseParameters message];
+  parameters.size = kSimpleDataLength;
+  [request.responseParametersArray addObject:parameters];
+  request.payload.body = [NSMutableData dataWithLength:kSimpleDataLength];
+
+  GRPCRequestOptions *callRequest =
+      [[GRPCRequestOptions alloc] initWithHost:(NSString *)kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = GRPCTransportTypeInsecure;
+  options.enableFlowControl = YES;
+  GRPCCall2 *call =
+      [[GRPCCall2 alloc] initWithRequestOptions:callRequest
+                                responseHandler:[[ClientTestsBlockCallbacks alloc]
+                                                    initWithInitialMetadataCallback:nil
+                                                                    messageCallback:nil
+                                                                      closeCallback:nil
+                                                                  writeDataCallback:^{
+                                                                    [expectWriteData fulfill];
+                                                                  }]
+                                    callOptions:options];
+
+  [call start];
+  [call receiveNextMessage];
+  [call writeData:[request data]];
+
+  // Wait for 3 seconds and make sure we do not receive the response
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+
+  [call finish];
+}
+
+- (void)testReadFlowControl {
+  __weak __block XCTestExpectation *expectBlockedMessage =
+      [self expectationWithDescription:@"Message not delivered without recvNextMessage"];
+  __weak __block XCTestExpectation *expectPassedMessage = nil;
+  __weak __block XCTestExpectation *expectBlockedClose =
+      [self expectationWithDescription:@"Call not closed with pending message"];
+  __weak __block XCTestExpectation *expectPassedClose = nil;
+  expectBlockedMessage.inverted = YES;
+  expectBlockedClose.inverted = YES;
+
+  RMTStreamingOutputCallRequest *request = [RMTStreamingOutputCallRequest message];
+  RMTResponseParameters *parameters = [RMTResponseParameters message];
+  parameters.size = kSimpleDataLength;
+  [request.responseParametersArray addObject:parameters];
+  request.payload.body = [NSMutableData dataWithLength:kSimpleDataLength];
+
+  GRPCRequestOptions *callRequest =
+      [[GRPCRequestOptions alloc] initWithHost:(NSString *)kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = GRPCTransportTypeInsecure;
+  options.enableFlowControl = YES;
+  __block int unblocked = NO;
+  GRPCCall2 *call = [[GRPCCall2 alloc]
+      initWithRequestOptions:callRequest
+             responseHandler:[[ClientTestsBlockCallbacks alloc] initWithInitialMetadataCallback:nil
+                                 messageCallback:^(NSData *message) {
+                                   if (!unblocked) {
+                                     [expectBlockedMessage fulfill];
+                                   } else {
+                                     [expectPassedMessage fulfill];
+                                   }
+                                 }
+                                 closeCallback:^(NSDictionary *trailers, NSError *error) {
+                                   if (!unblocked) {
+                                     [expectBlockedClose fulfill];
+                                   } else {
+                                     [expectPassedClose fulfill];
+                                   }
+                                 }]
+                 callOptions:options];
+
+  [call start];
+  [call writeData:[request data]];
+  [call finish];
+
+  // Wait to make sure we do not receive the response
+  [self waitForExpectationsWithTimeout:kInvertedTimeout handler:nil];
+
+  expectPassedMessage =
+      [self expectationWithDescription:@"Message delivered with receiveNextMessage"];
+  expectPassedClose = [self expectationWithDescription:@"Close delivered after receiveNextMessage"];
+
+  unblocked = YES;
+  [call receiveNextMessage];
+
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+}
+
+- (void)testReadFlowControlReadyBeforeStart {
+  __weak XCTestExpectation *expectBlockedMessage =
+      [self expectationWithDescription:@"Message delivered with receiveNextMessage"];
+  expectBlockedMessage.inverted = YES;
+
+  RMTStreamingOutputCallRequest *request = [RMTStreamingOutputCallRequest message];
+  RMTResponseParameters *parameters = [RMTResponseParameters message];
+  parameters.size = kSimpleDataLength;
+  [request.responseParametersArray addObject:parameters];
+  request.payload.body = [NSMutableData dataWithLength:kSimpleDataLength];
+
+  GRPCRequestOptions *callRequest =
+      [[GRPCRequestOptions alloc] initWithHost:(NSString *)kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = GRPCTransportTypeInsecure;
+  options.enableFlowControl = YES;
+  __block BOOL closed = NO;
+  GRPCCall2 *call = [[GRPCCall2 alloc]
+      initWithRequestOptions:callRequest
+             responseHandler:[[ClientTestsBlockCallbacks alloc]
+                                 initWithInitialMetadataCallback:nil
+                                                 messageCallback:^(NSData *message) {
+                                                   [expectBlockedMessage fulfill];
+                                                   XCTAssertFalse(closed);
+                                                 }
+                                                   closeCallback:nil]
+                 callOptions:options];
+
+  [call receiveNextMessage];
+  [call start];
+  [call writeData:[request data]];
+  [call finish];
+
+  [self waitForExpectationsWithTimeout:kInvertedTimeout handler:nil];
+}
+
+- (void)testReadFlowControlReadyAfterStart {
+  __weak XCTestExpectation *expectPassedMessage =
+      [self expectationWithDescription:@"Message delivered with receiveNextMessage"];
+  __weak XCTestExpectation *expectPassedClose =
+      [self expectationWithDescription:@"Close delivered with receiveNextMessage"];
+
+  RMTStreamingOutputCallRequest *request = [RMTStreamingOutputCallRequest message];
+  RMTResponseParameters *parameters = [RMTResponseParameters message];
+  parameters.size = kSimpleDataLength;
+  [request.responseParametersArray addObject:parameters];
+  request.payload.body = [NSMutableData dataWithLength:kSimpleDataLength];
+
+  GRPCRequestOptions *callRequest =
+      [[GRPCRequestOptions alloc] initWithHost:(NSString *)kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = GRPCTransportTypeInsecure;
+  options.enableFlowControl = YES;
+  __block BOOL closed = NO;
+  GRPCCall2 *call = [[GRPCCall2 alloc]
+      initWithRequestOptions:callRequest
+             responseHandler:[[ClientTestsBlockCallbacks alloc] initWithInitialMetadataCallback:nil
+                                 messageCallback:^(NSData *message) {
+                                   [expectPassedMessage fulfill];
+                                   XCTAssertFalse(closed);
+                                 }
+                                 closeCallback:^(NSDictionary *trailers, NSError *error) {
+                                   closed = YES;
+                                   [expectPassedClose fulfill];
+                                 }]
+                 callOptions:options];
+
+  [call start];
+  [call receiveNextMessage];
+  [call writeData:[request data]];
+  [call finish];
+
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+}
+
+- (void)testReadFlowControlNonBlockingFailure {
+  __weak XCTestExpectation *completion = [self expectationWithDescription:@"RPC completed."];
+
+  GRPCRequestOptions *requestOptions =
+      [[GRPCRequestOptions alloc] initWithHost:kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.enableFlowControl = YES;
+  options.transportType = GRPCTransportTypeInsecure;
+
+  RMTSimpleRequest *request = [RMTSimpleRequest message];
+  request.payload.body = [NSMutableData dataWithLength:options.responseSizeLimit];
+
+  RMTEchoStatus *status = [RMTEchoStatus message];
+  status.code = 2;
+  status.message = @"test";
+  request.responseStatus = status;
+
+  GRPCCall2 *call = [[GRPCCall2 alloc]
+      initWithRequestOptions:requestOptions
+             responseHandler:[[ClientTestsBlockCallbacks alloc] initWithInitialMetadataCallback:nil
+                                 messageCallback:^(NSData *data) {
+                                   XCTFail(@"Received unexpected message");
+                                 }
+                                 closeCallback:^(NSDictionary *trailingMetadata, NSError *error) {
+                                   XCTAssertNotNil(error, @"Expecting non-nil error");
+                                   XCTAssertEqual(error.code, 2);
+                                   [completion fulfill];
+                                 }]
+                 callOptions:options];
+  [call writeData:[request data]];
+  [call start];
+  [call finish];
+
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+}
+
 @end

src/objective-c/tests/InteropTests.m

@@ -77,6 +77,11 @@ BOOL isRemoteInteropTest(NSString *host) {
 // Convenience class to use blocks as callbacks
 @interface InteropTestsBlockCallbacks : NSObject<GRPCProtoResponseHandler>
 
+- (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
+                                messageCallback:(void (^)(id))messageCallback
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback
+                           writeMessageCallback:(void (^)(void))writeMessageCallback;
+
 - (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
                                 messageCallback:(void (^)(id))messageCallback
                                   closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback;
@@ -87,21 +92,33 @@ BOOL isRemoteInteropTest(NSString *host) {
   void (^_initialMetadataCallback)(NSDictionary *);
   void (^_messageCallback)(id);
   void (^_closeCallback)(NSDictionary *, NSError *);
+  void (^_writeMessageCallback)(void);
   dispatch_queue_t _dispatchQueue;
 }
 
 - (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
                                 messageCallback:(void (^)(id))messageCallback
-                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback {
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback
+                           writeMessageCallback:(void (^)(void))writeMessageCallback {
   if ((self = [super init])) {
     _initialMetadataCallback = initialMetadataCallback;
     _messageCallback = messageCallback;
     _closeCallback = closeCallback;
+    _writeMessageCallback = writeMessageCallback;
     _dispatchQueue = dispatch_queue_create(nil, DISPATCH_QUEUE_SERIAL);
   }
   return self;
 }
 
+- (instancetype)initWithInitialMetadataCallback:(void (^)(NSDictionary *))initialMetadataCallback
+                                messageCallback:(void (^)(id))messageCallback
+                                  closeCallback:(void (^)(NSDictionary *, NSError *))closeCallback {
+  return [self initWithInitialMetadataCallback:initialMetadataCallback
+                               messageCallback:messageCallback
+                                 closeCallback:closeCallback
+                          writeMessageCallback:nil];
+}
+
 - (void)didReceiveInitialMetadata:(NSDictionary *)initialMetadata {
   if (_initialMetadataCallback) {
     _initialMetadataCallback(initialMetadata);
@@ -120,6 +137,12 @@ BOOL isRemoteInteropTest(NSString *host) {
   }
 }
 
+- (void)didWriteMessage {
+  if (_writeMessageCallback) {
+    _writeMessageCallback();
+  }
+}
+
 - (dispatch_queue_t)dispatchQueue {
   return _dispatchQueue;
 }
@@ -570,6 +593,67 @@ BOOL isRemoteInteropTest(NSString *host) {
   [self waitForExpectationsWithTimeout:TEST_TIMEOUT handler:nil];
 }
 
+- (void)testPingPongRPCWithFlowControl {
+  XCTAssertNotNil([[self class] host]);
+  __weak XCTestExpectation *expectation = [self expectationWithDescription:@"PingPongWithV2API"];
+
+  NSArray *requests = @[ @27182, @8, @1828, @45904 ];
+  NSArray *responses = @[ @31415, @9, @2653, @58979 ];
+
+  __block int index = 0;
+
+  id request = [RMTStreamingOutputCallRequest messageWithPayloadSize:requests[index]
+                                               requestedResponseSize:responses[index]];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = [[self class] transportType];
+  options.PEMRootCertificates = [[self class] PEMRootCertificates];
+  options.hostNameOverride = [[self class] hostNameOverride];
+  options.enableFlowControl = YES;
+  __block BOOL canWriteData = NO;
+
+  __block GRPCStreamingProtoCall *call = [_service
+      fullDuplexCallWithResponseHandler:[[InteropTestsBlockCallbacks alloc]
+                                            initWithInitialMetadataCallback:nil
+                                            messageCallback:^(id message) {
+                                              XCTAssertLessThan(index, 4,
+                                                                @"More than 4 responses received.");
+                                              id expected = [RMTStreamingOutputCallResponse
+                                                  messageWithPayloadSize:responses[index]];
+                                              XCTAssertEqualObjects(message, expected);
+                                              index += 1;
+                                              if (index < 4) {
+                                                id request = [RMTStreamingOutputCallRequest
+                                                    messageWithPayloadSize:requests[index]
+                                                     requestedResponseSize:responses[index]];
+                                                XCTAssertTrue(canWriteData);
+                                                canWriteData = NO;
+                                                [call writeMessage:request];
+                                                [call receiveNextMessage];
+                                              } else {
+                                                [call finish];
+                                              }
+                                            }
+                                            closeCallback:^(NSDictionary *trailingMetadata,
+                                                            NSError *error) {
+                                              XCTAssertNil(error,
+                                                           @"Finished with unexpected error: %@",
+                                                           error);
+                                              XCTAssertEqual(index, 4,
+                                                             @"Received %i responses instead of 4.",
+                                                             index);
+                                              [expectation fulfill];
+                                            }
+                                            writeMessageCallback:^{
+                                              canWriteData = YES;
+                                            }]
+                            callOptions:options];
+  [call start];
+  [call receiveNextMessage];
+  [call writeMessage:request];
+
+  [self waitForExpectationsWithTimeout:TEST_TIMEOUT handler:nil];
+}
+
 - (void)testEmptyStreamRPC {
   XCTAssertNotNil([[self class] host]);
   __weak XCTestExpectation *expectation = [self expectationWithDescription:@"EmptyStream"];