[ReactNative] Add profiling hooks to bridge modules at runtime

Summary:
@public

Add `RCT_DEV`-only profiling hooks to every method in every bridge modules to
add a little bit more detail to the timeline profile.

Test Plan: {F22522834}

React/Base/RCTBridge.m

@@ -1313,7 +1313,12 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
 
 #pragma mark - Payload Generation
 
-- (void)dispatchBlock:(dispatch_block_t)block forModule:(NSNumber *)moduleID
+- (void)dispatchBlock:(dispatch_block_t)block forModule:(id<RCTBridgeModule>)module
+{
+  [self dispatchBlock:block forModuleID:RCTModuleIDsByName[RCTBridgeModuleNameForClass([module class])]];
+}
+
+- (void)dispatchBlock:(dispatch_block_t)block forModuleID:(NSNumber *)moduleID
 {
   RCTAssertJSThread();
 
@@ -1458,7 +1463,7 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
     if ([module respondsToSelector:@selector(batchDidComplete)]) {
       [self dispatchBlock:^{
         [module batchDidComplete];
-      } forModule:moduleID];
+      } forModuleID:moduleID];
     }
   }];
 }
@@ -1526,7 +1531,7 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
       @"selector": NSStringFromSelector(method.selector),
       @"args": RCTJSONStringify(params ?: [NSNull null], NULL),
     });
-  } forModule:@(moduleID)];
+  } forModuleID:@(moduleID)];
 
   return YES;
 }
@@ -1546,7 +1551,7 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
         RCTProfileBeginEvent();
         [observer didUpdateFrame:frameUpdate];
         RCTProfileEndEvent(name, @"objc_call,fps", nil);
-      } forModule:RCTModuleIDsByName[RCTBridgeModuleNameForClass([observer class])]];
+      } forModule:(id<RCTBridgeModule>)observer];
     }
   }
 
@@ -1598,14 +1603,17 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
     return;
   }
 
-  RCTProfileInit();
+  [_javaScriptExecutor executeBlockOnJavaScriptQueue:^{
+    RCTProfileInit(self);
+  }];
 }
 
 - (void)stopProfiling
 {
    RCTAssertMainThread();
 
-  NSString *log = RCTProfileEnd();
+  [_javaScriptExecutor executeBlockOnJavaScriptQueue:^{
+    NSString *log = RCTProfileEnd(self);
     NSURL *bundleURL = _parentBridge.bundleURL;
     NSString *URLString = [NSString stringWithFormat:@"%@://%@:%@/profile", bundleURL.scheme, bundleURL.host, bundleURL.port];
     NSURL *URL = [NSURL URLWithString:URLString];
@@ -1620,6 +1628,7 @@ RCT_INNER_BRIDGE_ONLY(_invokeAndProcessModule:(NSString *)module method:(NSStrin
                                                                  }
                                                                }];
     [task resume];
+  }];
 }
 
 @end

React/Base/RCTProfile.h

@@ -25,6 +25,8 @@ NSString *const RCTProfileDidEndProfiling;
 
 #if RCT_DEV
 
+@class RCTBridge;
+
 #define RCTProfileBeginFlowEvent() \
 _Pragma("clang diagnostic push") \
 _Pragma("clang diagnostic ignored \"-Wshadow\"") \
@@ -45,14 +47,14 @@ RCT_EXTERN BOOL RCTProfileIsProfiling(void);
 /**
  * Start collecting profiling information
  */
-RCT_EXTERN void RCTProfileInit(void);
+RCT_EXTERN void RCTProfileInit(RCTBridge *);
 
 /**
  * Stop profiling and return a JSON string of the collected data - The data
  * returned is compliant with google's trace event format - the format used
  * as input to trace-viewer
  */
-RCT_EXTERN NSString *RCTProfileEnd(void);
+RCT_EXTERN NSString *RCTProfileEnd(RCTBridge *);
 
 /**
  * Collects the initial event information for the event and returns a reference ID

React/Base/RCTProfile.m

@@ -10,10 +10,13 @@
 #import "RCTProfile.h"
 
 #import <mach/mach.h>
+#import <objc/message.h>
+#import <objc/runtime.h>
 
 #import <UIKit/UIKit.h>
 
 #import "RCTAssert.h"
+#import "RCTBridge.h"
 #import "RCTDefines.h"
 #import "RCTUtils.h"
 
@@ -32,6 +35,7 @@ NSDictionary *RCTProfileGetMemoryUsage(void);
 
 NSString const *RCTProfileTraceEvents = @"traceEvents";
 NSString const *RCTProfileSamples = @"samples";
+NSString *const RCTProfilePrefix = @"rct_profile_";
 
 #pragma mark - Variables
 
@@ -92,6 +96,111 @@ NSDictionary *RCTProfileGetMemoryUsage(void)
   }
 }
 
+#pragma mark - Module hooks
+
+@interface RCTBridge (Private)
+
+- (void)dispatchBlock:(dispatch_block_t)block forModule:(id<RCTBridgeModule>)module;
+
+@end
+
+static const char *RCTProfileProxyClassName(Class);
+static const char *RCTProfileProxyClassName(Class class)
+{
+  return [RCTProfilePrefix stringByAppendingString:NSStringFromClass(class)].UTF8String;
+}
+
+static SEL RCTProfileProxySelector(SEL);
+static SEL RCTProfileProxySelector(SEL selector)
+{
+  NSString *selectorName = NSStringFromSelector(selector);
+  return NSSelectorFromString([RCTProfilePrefix stringByAppendingString:selectorName]);
+}
+
+static void RCTProfileForwardInvocation(NSObject *, SEL, NSInvocation *);
+static void RCTProfileForwardInvocation(NSObject *self, SEL cmd, NSInvocation *invocation)
+{
+  NSString *name = [NSString stringWithFormat:@"-[%@ %@]", NSStringFromClass([self class]), NSStringFromSelector(invocation.selector)];
+  SEL newSel = RCTProfileProxySelector(invocation.selector);
+
+  if ([object_getClass(self) instancesRespondToSelector:newSel]) {
+    invocation.selector = newSel;
+    RCTProfileBeginEvent();
+    [invocation invoke];
+    RCTProfileEndEvent(name, @"objc_call,modules,auto", nil);
+  } else {
+    // Use original selector to don't change error message
+    [self doesNotRecognizeSelector:invocation.selector];
+  }
+}
+
+static IMP RCTProfileMsgForward(NSObject *, SEL);
+static IMP RCTProfileMsgForward(NSObject *self, SEL selector)
+{
+  IMP imp = _objc_msgForward;
+#if !defined(__arm64__)
+  NSMethodSignature *signature = [self methodSignatureForSelector:selector];
+  if (signature.methodReturnType[0] == _C_STRUCT_B && signature.methodReturnLength > 8) {
+    imp = _objc_msgForward_stret;
+  }
+#endif
+  return imp;
+}
+
+static void RCTProfileHookModules(RCTBridge *);
+static void RCTProfileHookModules(RCTBridge *bridge)
+{
+  [bridge.modules enumerateKeysAndObjectsUsingBlock:^(NSString *className, id<RCTBridgeModule> module, BOOL *stop) {
+    [bridge dispatchBlock:^{
+      Class moduleClass = object_getClass(module);
+      Class proxyClass = objc_allocateClassPair(moduleClass, RCTProfileProxyClassName(moduleClass), 0);
+
+      unsigned int methodCount;
+      Method *methods = class_copyMethodList(moduleClass, &methodCount);
+      for (NSUInteger i = 0; i < methodCount; i++) {
+        Method method = methods[i];
+        SEL selector = method_getName(method);
+        if ([NSStringFromSelector(selector) hasPrefix:@"rct"] || [NSObject instancesRespondToSelector:selector]) {
+          continue;
+        }
+        IMP originalIMP = method_getImplementation(method);
+        const char *returnType = method_getTypeEncoding(method);
+        class_addMethod(proxyClass, selector, RCTProfileMsgForward(module, selector), returnType);
+        class_addMethod(proxyClass, RCTProfileProxySelector(selector), originalIMP, returnType);
+      }
+      free(methods);
+
+      for (Class cls in @[proxyClass, object_getClass(proxyClass)]) {
+        Method oldImp = class_getInstanceMethod(cls, @selector(class));
+        class_replaceMethod(cls, @selector(class), imp_implementationWithBlock(^{ return moduleClass; }), method_getTypeEncoding(oldImp));
+      }
+
+      IMP originalFwd = class_replaceMethod(moduleClass, @selector(forwardInvocation:), (IMP)RCTProfileForwardInvocation, "v@:@");
+      if (originalFwd != NULL) {
+        class_addMethod(proxyClass, RCTProfileProxySelector(@selector(forwardInvocation:)), originalFwd, "v@:@");
+      }
+
+      objc_registerClassPair(proxyClass);
+      object_setClass(module, proxyClass);
+    } forModule:module];
+  }];
+}
+
+void RCTProfileUnhookModules(RCTBridge *);
+void RCTProfileUnhookModules(RCTBridge *bridge)
+{
+  [bridge.modules enumerateKeysAndObjectsUsingBlock:^(NSString *className, id<RCTBridgeModule> module, BOOL *stop) {
+    [bridge dispatchBlock:^{
+      Class proxyClass = object_getClass(module);
+      if (module.class != proxyClass) {
+        object_setClass(module, module.class);
+        objc_disposeClassPair(proxyClass);
+      }
+    } forModule:module];
+  }];
+}
+
+
 #pragma mark - Public Functions
 
 BOOL RCTProfileIsProfiling(void)
@@ -102,8 +211,10 @@ BOOL RCTProfileIsProfiling(void)
   return profiling;
 }
 
-void RCTProfileInit(void)
+void RCTProfileInit(RCTBridge *bridge)
 {
+  RCTProfileHookModules(bridge);
+
   static dispatch_once_t onceToken;
   dispatch_once(&onceToken, ^{
     _RCTProfileLock = [[NSLock alloc] init];
@@ -121,7 +232,7 @@ void RCTProfileInit(void)
                                                       object:nil];
 }
 
-NSString *RCTProfileEnd(void)
+NSString *RCTProfileEnd(RCTBridge *bridge)
 {
   [[NSNotificationCenter defaultCenter] postNotificationName:RCTProfileDidEndProfiling
                                                       object:nil];
@@ -132,6 +243,9 @@ NSString *RCTProfileEnd(void)
     RCTProfileInfo = nil;
     RCTProfileOngoingEvents = nil;
   );
+
+  RCTProfileUnhookModules(bridge);
+
   return log;
 }