Merge pull request #17936 from muxi/flow-control

gRPC Objective-C Flow Control

src/objective-c/GRPCClient/GRPCCall.h

@@ -184,6 +184,12 @@ extern NSString *const kGRPCTrailersKey;
 - (void)didCloseWithTrailingMetadata:(nullable NSDictionary *)trailingMetadata
                                error:(nullable NSError *)error;
 
+/**
+ * Issued when flow control is enabled for the call and a message written with writeData: method of
+ * GRPCCall2 is passed to gRPC core with SEND_MESSAGE operation.
+ */
+- (void)didWriteData;
+
 @end
 
 /**
@@ -264,6 +270,14 @@ extern NSString *const kGRPCTrailersKey;
  */
 - (void)finish;
 
+/**
+ * Tell gRPC to receive the next N gRPC message from gRPC core.
+ *
+ * This method should only be used when flow control is enabled. When flow control is not enabled,
+ * this method is a no-op.
+ */
+- (void)receiveNextMessages:(NSUInteger)numberOfMessages;
+
 /**
  * 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)receiveNextMessages:(NSUInteger)numberOfMessages;
+
 @end
 
 @implementation GRPCRequestOptions
@@ -113,6 +122,8 @@ const char *kCFStreamVarName = "grpc_cfstream";
   BOOL _canceled;
   /** Flags whether call has been finished. */
   BOOL _finished;
+  /** The number of pending messages receiving requests. */
+  NSUInteger _pendingReceiveNextMessages;
 }
 
 - (instancetype)initWithRequestOptions:(GRPCRequestOptions *)requestOptions
@@ -190,11 +201,22 @@ 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];
+                                     }
+                                   }
+                                 }];
     [_call setResponseDispatchQueue:_dispatchQueue];
     if (_callOptions.initialMetadata) {
       [_call.requestHeaders addEntriesFromDictionary:_callOptions.initialMetadata];
     }
+    if (_pendingReceiveNextMessages > 0) {
+      [_call receiveNextMessages:_pendingReceiveNextMessages];
+      _pendingReceiveNextMessages = 0;
+    }
     copiedCall = _call;
   }
 
@@ -364,6 +386,33 @@ const char *kCFStreamVarName = "grpc_cfstream";
   }
 }
 
+- (void)issueDidWriteData {
+  @synchronized(self) {
+    if (_callOptions.flowControlEnabled && [_handler respondsToSelector:@selector(didWriteData)]) {
+      dispatch_async(_dispatchQueue, ^{
+        id<GRPCResponseHandler> copiedHandler = nil;
+        @synchronized(self) {
+          copiedHandler = self->_handler;
+        };
+        [copiedHandler didWriteData];
+      });
+    }
+  }
+}
+
+- (void)receiveNextMessages:(NSUInteger)numberOfMessages {
+  // branching based on _callOptions.flowControlEnabled is handled inside _call
+  GRPCCall *copiedCall = nil;
+  @synchronized(self) {
+    copiedCall = _call;
+    if (copiedCall == nil) {
+      _pendingReceiveNextMessages += numberOfMessages;
+      return;
+    }
+  }
+  [copiedCall receiveNextMessages:numberOfMessages];
+}
+
 @end
 
 // The following methods of a C gRPC call object aren't reentrant, and thus
@@ -427,6 +476,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 pending read message request from user.
+  NSUInteger _pendingReceiveNextMessages;
 }
 
 @synthesize state = _state;
@@ -486,12 +544,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;
@@ -499,7 +571,7 @@ const char *kCFStreamVarName = "grpc_cfstream";
   if (safety > GRPCCallSafetyCacheableRequest) {
     return nil;
   }
-  if (requestWriter.state != GRXWriterStateNotStarted) {
+  if (requestsWriter.state != GRXWriterStateNotStarted) {
     return nil;
   }
 
@@ -512,16 +584,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
+    _pendingReceiveNextMessages = 0;
+    _pendingCoreRead = YES;
   }
   return self;
 }
@@ -593,11 +669,16 @@ 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.flowControlEnabled && (_pendingCoreRead || _pendingReceiveNextMessages == 0)) {
+      return;
+    }
+    _pendingCoreRead = YES;
+    _pendingReceiveNextMessages--;
   }
 
   dispatch_async(_callQueue, ^{
@@ -620,6 +701,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->_pendingCoreRead = NO;
           [strongSelf
               finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
                                                   code:GRPCErrorCodeResourceExhausted
@@ -635,7 +717,13 @@ const char *kCFStreamVarName = "grpc_cfstream";
         @synchronized(strongSelf) {
           [strongSelf->_responseWriteable enqueueValue:data
                                      completionHandler:^{
-                                       [strongSelf startNextRead];
+                                       __strong GRPCCall *strongSelf = weakSelf;
+                                       if (strongSelf) {
+                                         @synchronized(strongSelf) {
+                                           strongSelf->_pendingCoreRead = NO;
+                                           [strongSelf maybeStartNextRead];
+                                         }
+                                       }
                                      }];
         }
       }
@@ -686,6 +774,20 @@ const char *kCFStreamVarName = "grpc_cfstream";
   });
 }
 
+- (void)receiveNextMessages:(NSUInteger)numberOfMessages {
+  if (numberOfMessages == 0) {
+    return;
+  }
+  @synchronized(self) {
+    _pendingReceiveNextMessages += numberOfMessages;
+
+    if (_state != GRXWriterStateStarted || !_callOptions.flowControlEnabled) {
+      return;
+    }
+    [self maybeStartNextRead];
+  }
+}
+
 #pragma mark GRXWriteable implementation
 
 // Only called from the call queue. The error handler will be called from the
@@ -699,9 +801,11 @@ const char *kCFStreamVarName = "grpc_cfstream";
     GRPCCall *strongSelf = weakSelf;
     if (strongSelf) {
       strongSelf->_requestWriter.state = GRXWriterStateStarted;
+      if (strongSelf->_writeDone) {
+        strongSelf->_writeDone();
+      }
     }
   };
-
   GRPCOpSendMessage *op =
       [[GRPCOpSendMessage alloc] initWithMessage:message handler:resumingHandler];
   if (!_unaryCall) {
@@ -778,8 +882,11 @@ const char *kCFStreamVarName = "grpc_cfstream";
     // Response headers received.
     __strong GRPCCall *strongSelf = weakSelf;
     if (strongSelf) {
-      strongSelf.responseHeaders = headers;
-      [strongSelf startNextRead];
+      @synchronized(strongSelf) {
+        strongSelf.responseHeaders = headers;
+        strongSelf->_pendingCoreRead = NO;
+        [strongSelf maybeStartNextRead];
+      }
     }
   }
       completionHandler:^(NSError *error, NSDictionary *trailers) {
@@ -933,7 +1040,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 flowControlEnabled;
+
 // OAuth2 parameters. Users of gRPC may specify one of the following two parameters.
 
 /**
@@ -232,6 +240,19 @@ 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 can issue a didReceiveMessage
+ * callback.
+ *
+ * If writeData: method is called more than once before issuance of a didWriteData callback, gRPC
+ * will continue to queue the message and write them to gRPC core in order. However, the user
+ * assumes their own responsibility of flow control by keeping tracking of the pending writes in
+ * the call.
+ */
+@property(readwrite) BOOL flowControlEnabled;
+
 // 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 kDefaultFlowControlEnabled = 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 _flowControlEnabled;
   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 flowControlEnabled = _flowControlEnabled;
 @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
+                    flowControlEnabled:kDefaultFlowControlEnabled
                      oauth2AccessToken:kDefaultOauth2AccessToken
                      authTokenProvider:kDefaultAuthTokenProvider
                        initialMetadata:kDefaultInitialMetadata
@@ -134,6 +138,7 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
 
 - (instancetype)initWithServerAuthority:(NSString *)serverAuthority
                                 timeout:(NSTimeInterval)timeout
+                     flowControlEnabled:(BOOL)flowControlEnabled
                       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;
+    _flowControlEnabled = flowControlEnabled;
     _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
+                                                 flowControlEnabled:_flowControlEnabled
                                                   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
+           flowControlEnabled:_flowControlEnabled
             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 flowControlEnabled;
 @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
+                    flowControlEnabled:kDefaultFlowControlEnabled
                      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
+                                                 flowControlEnabled:_flowControlEnabled
                                                   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
+           flowControlEnabled:_flowControlEnabled
             oauth2AccessToken:_oauth2AccessToken
             authTokenProvider:_authTokenProvider
               initialMetadata:_initialMetadata
@@ -417,6 +429,10 @@ static BOOL areObjectsEqual(id obj1, id obj2) {
   }
 }
 
+- (void)setFlowControlEnabled:(BOOL)flowControlEnabled {
+  _flowControlEnabled = flowControlEnabled;
+}
+
 - (void)setOauth2AccessToken:(NSString *)oauth2AccessToken {
   _oauth2AccessToken = [oauth2AccessToken copy];
 }

src/objective-c/ProtoRPC/ProtoRPC.h

@@ -57,6 +57,13 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)didCloseWithTrailingMetadata:(nullable NSDictionary *)trailingMetadata
                                error:(nullable NSError *)error;
 
+/**
+ * Issued when flow control is enabled for the call and a message (written with writeMessage: method
+ * of GRPCStreamingProtoCall or the initializer of GRPCUnaryProtoCall) is passed to gRPC core with
+ * SEND_MESSAGE operation.
+ */
+- (void)didWriteMessage;
+
 @end
 
 /** A unary-request RPC call with Protobuf. */
@@ -130,6 +137,26 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)finish;
 
+/**
+ * Tell gRPC to receive another message.
+ *
+ * This method should only be used when flow control is enabled. If flow control is enabled, gRPC
+ * will only receive additional messages after the user indicates so by using either
+ * receiveNextMessage: or receiveNextMessages: methods. If flow control is not enabled, messages
+ * will be automatically received after the previous one is delivered.
+ */
+- (void)receiveNextMessage;
+
+/**
+ * Tell gRPC to receive another N message.
+ *
+ * This method should only be used when flow control is enabled. If flow control is enabled, the
+ * messages received from the server are buffered in gRPC until the user want to receive the next
+ * message. If flow control is not enabled, messages will be automatically received after the
+ * previous one is delivered.
+ */
+- (void)receiveNextMessages:(NSUInteger)numberOfMessages;
+
 @end
 
 NS_ASSUME_NONNULL_END

src/objective-c/ProtoRPC/ProtoRPC.m

@@ -71,6 +71,7 @@ static NSError *ErrorForBadProto(id proto, Class expectedClass, NSError *parsing
 
 - (void)start {
   [_call start];
+  [_call receiveNextMessage];
   [_call writeMessage:_message];
   [_call finish];
 }
@@ -196,6 +197,17 @@ static NSError *ErrorForBadProto(id proto, Class expectedClass, NSError *parsing
   [copiedCall finish];
 }
 
+- (void)receiveNextMessage {
+  [self receiveNextMessages:1];
+}
+- (void)receiveNextMessages:(NSUInteger)numberOfMessages {
+  GRPCCall2 *copiedCall;
+  @synchronized(self) {
+    copiedCall = _call;
+  }
+  [copiedCall receiveNextMessages:numberOfMessages];
+}
+
 - (void)didReceiveInitialMetadata:(NSDictionary *)initialMetadata {
   @synchronized(self) {
     if (initialMetadata != nil &&
@@ -259,6 +271,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

@@ -40,11 +40,13 @@ static NSString *const kService = @"TestService";
 static GRPCProtoMethod *kInexistentMethod;
 static GRPCProtoMethod *kEmptyCallMethod;
 static GRPCProtoMethod *kUnaryCallMethod;
+static GRPCProtoMethod *kOutputStreamingCallMethod;
 static GRPCProtoMethod *kFullDuplexCallMethod;
 
 static const int kSimpleDataLength = 100;
 
-static const NSTimeInterval kTestTimeout = 16;
+static const NSTimeInterval kTestTimeout = 8;
+static const NSTimeInterval kInvertedTimeout = 2;
 
 // Reveal the _class ivar for testing access
 @interface GRPCCall2 () {
@@ -57,6 +59,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;
@@ -67,21 +74,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);
@@ -100,6 +119,12 @@ static const NSTimeInterval kTestTimeout = 16;
   }
 }
 
+- (void)didWriteData {
+  if (self->_writeDataCallback) {
+    self->_writeDataCallback();
+  }
+}
+
 - (dispatch_queue_t)dispatchQueue {
   return _dispatchQueue;
 }
@@ -120,6 +145,9 @@ static const NSTimeInterval kTestTimeout = 16;
       [[GRPCProtoMethod alloc] initWithPackage:kPackage service:kService method:@"EmptyCall"];
   kUnaryCallMethod =
       [[GRPCProtoMethod alloc] initWithPackage:kPackage service:kService method:@"UnaryCall"];
+  kOutputStreamingCallMethod = [[GRPCProtoMethod alloc] initWithPackage:kPackage
+                                                                service:kService
+                                                                 method:@"StreamingOutputCall"];
   kFullDuplexCallMethod =
       [[GRPCProtoMethod alloc] initWithPackage:kPackage service:kService method:@"FullDuplexCall"];
 }
@@ -478,4 +506,268 @@ static const NSTimeInterval kTestTimeout = 16;
   [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
 }
 
+- (void)testFlowControlWrite {
+  __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.flowControlEnabled = YES;
+  GRPCCall2 *call =
+      [[GRPCCall2 alloc] initWithRequestOptions:callRequest
+                                responseHandler:[[ClientTestsBlockCallbacks alloc]
+                                                    initWithInitialMetadataCallback:nil
+                                                                    messageCallback:nil
+                                                                      closeCallback:nil
+                                                                  writeDataCallback:^{
+                                                                    [expectWriteData fulfill];
+                                                                  }]
+                                    callOptions:options];
+
+  [call start];
+  [call receiveNextMessages:1];
+  [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)testFlowControlRead {
+  __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;
+
+  RMTSimpleRequest *request = [RMTSimpleRequest message];
+  request.responseSize = kSimpleDataLength;
+  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.flowControlEnabled = 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 receiveNextMessages:1];
+
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+}
+
+- (void)testFlowControlMultipleMessages {
+  __weak XCTestExpectation *expectPassedMessage =
+      [self expectationWithDescription:@"two messages delivered with receiveNextMessage"];
+  expectPassedMessage.expectedFulfillmentCount = 2;
+  __weak XCTestExpectation *expectBlockedMessage =
+      [self expectationWithDescription:@"Message 3 not delivered"];
+  expectBlockedMessage.inverted = YES;
+  __weak XCTestExpectation *expectWriteTwice =
+      [self expectationWithDescription:@"Write 2 messages done"];
+  expectWriteTwice.expectedFulfillmentCount = 2;
+
+  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:kFullDuplexCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.transportType = GRPCTransportTypeInsecure;
+  options.flowControlEnabled = YES;
+  __block NSUInteger messageId = 0;
+  __block GRPCCall2 *call = [[GRPCCall2 alloc]
+      initWithRequestOptions:callRequest
+             responseHandler:[[ClientTestsBlockCallbacks alloc] initWithInitialMetadataCallback:nil
+                                 messageCallback:^(NSData *message) {
+                                   if (messageId <= 1) {
+                                     [expectPassedMessage fulfill];
+                                   } else {
+                                     [expectBlockedMessage fulfill];
+                                   }
+                                   messageId++;
+                                 }
+                                 closeCallback:nil
+                                 writeDataCallback:^{
+                                   [expectWriteTwice fulfill];
+                                 }]
+                 callOptions:options];
+
+  [call receiveNextMessages:2];
+  [call start];
+  [call writeData:[request data]];
+  [call writeData:[request data]];
+
+  [self waitForExpectationsWithTimeout:kInvertedTimeout handler:nil];
+}
+
+- (void)testFlowControlReadReadyBeforeStart {
+  __weak XCTestExpectation *expectPassedMessage =
+      [self expectationWithDescription:@"Message delivered with receiveNextMessage"];
+  __weak XCTestExpectation *expectPassedClose =
+      [self expectationWithDescription:@"Close delivered with receiveNextMessage"];
+
+  RMTSimpleRequest *request = [RMTSimpleRequest message];
+  request.responseSize = kSimpleDataLength;
+  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.flowControlEnabled = 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 *ttrailers, NSError *error) {
+                                   closed = YES;
+                                   [expectPassedClose fulfill];
+                                 }]
+                 callOptions:options];
+
+  [call receiveNextMessages:1];
+  [call start];
+  [call writeData:[request data]];
+  [call finish];
+
+  [self waitForExpectationsWithTimeout:kInvertedTimeout handler:nil];
+}
+
+- (void)testFlowControlReadReadyAfterStart {
+  __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.flowControlEnabled = 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 receiveNextMessages:1];
+  [call writeData:[request data]];
+  [call finish];
+
+  [self waitForExpectationsWithTimeout:kTestTimeout handler:nil];
+}
+
+- (void)testFlowControlReadNonBlockingFailure {
+  __weak XCTestExpectation *completion = [self expectationWithDescription:@"RPC completed."];
+
+  GRPCRequestOptions *requestOptions =
+      [[GRPCRequestOptions alloc] initWithHost:kHostAddress
+                                          path:kUnaryCallMethod.HTTPPath
+                                        safety:GRPCCallSafetyDefault];
+  GRPCMutableCallOptions *options = [[GRPCMutableCallOptions alloc] init];
+  options.flowControlEnabled = 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

@@ -79,6 +79,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;
@@ -89,21 +94,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);
@@ -122,6 +139,12 @@ BOOL isRemoteInteropTest(NSString *host) {
   }
 }
 
+- (void)didWriteMessage {
+  if (_writeMessageCallback) {
+    _writeMessageCallback();
+  }
+}
+
 - (dispatch_queue_t)dispatchQueue {
   return _dispatchQueue;
 }
@@ -702,6 +725,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.flowControlEnabled = 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"];