feat: extend middleware for session replay

amplitude · Aug 9, 2024 · 15ffe0a · 15ffe0a
1 parent b7a66d2
commit 15ffe0a
Show file tree

Hide file tree

Showing 8 changed files with 220 additions and 10 deletions.
diff --git a/Sources/Amplitude/AMPMiddleware.m b/Sources/Amplitude/AMPMiddleware.m
@@ -38,6 +38,12 @@ - (instancetype _Nonnull)initWithEvent:(NSMutableDictionary *_Nonnull) event wit
 
 @implementation AMPBlockMiddleware
 
+- (instancetype)init {
+    return [self initWithBlock:^(AMPMiddlewarePayload *payload, AMPMiddlewareNext next) {
+        next(payload);
+    }];
+}
+
 - (instancetype _Nonnull)initWithBlock:(AMPMiddlewareBlock)block {
     if (self = [super init]) {
         _block = block;
@@ -49,4 +55,16 @@ - (void)run:(AMPMiddlewarePayload *)payload next:(AMPMiddlewareNext)next {
     self.block(payload, next);
 }
 
+- (void)amplitudeDidFinishInitializing:(Amplitude *)amplitude {
+    if (self.didFinishInitializing) {
+        self.didFinishInitializing(amplitude);
+    }
+}
+
+- (void)amplitude:(Amplitude *)amplitude didUploadEventsManually:(BOOL)manually {
+    if (self.didUploadEventsManually) {
+        self.didUploadEventsManually(amplitude, manually);
+    }
+}
+
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.h b/Sources/Amplitude/AMPMiddlewareRunner.h
@@ -23,6 +23,8 @@
 #import <Foundation/Foundation.h>
 #import "AMPMiddleware.h"
 
+@class Amplitude;
+
 @interface AMPMiddlewareRunner : NSObject
 
 @property (nonatomic, nonnull, readonly) NSMutableArray<id<AMPMiddleware>> *middlewares;
@@ -33,4 +35,11 @@
 
 - (void) run:(AMPMiddlewarePayload *_Nonnull)payload next:(AMPMiddlewareNext _Nonnull)next;
 
+- (void)dispatchAmplitudeInitialized:(nonnull Amplitude *)amplitude;
+- (void)dispatchAmplitudeInitialized:(nonnull Amplitude *)amplitude
+                        toMiddleware:(nonnull id<AMPMiddleware>)middleware;
+
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude
+  didUploadEventsManually:(BOOL)isManualUpload;
+
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.m b/Sources/Amplitude/AMPMiddlewareRunner.m
@@ -22,6 +22,7 @@
 //
 
 #import <Foundation/Foundation.h>
+#import <objc/runtime.h>
 #import "AMPMiddlewareRunner.h"
 #import "AMPMiddleware.h"
 
@@ -62,4 +63,37 @@ - (void) runMiddlewares:(NSArray<id<AMPMiddleware>> *_Nonnull)middlewares
     }];
 }
 
+- (void)dispatchAmplitudeInitialized:(Amplitude *)amplitude {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        [self dispatchAmplitudeInitialized:amplitude toMiddleware:middleware];
+    }
+}
+
+- (void)dispatchAmplitudeInitialized:(Amplitude *)amplitude
+                        toMiddleware:(id<AMPMiddleware>)middleware {
+    if ([AMPMiddlewareRunner object:middleware
+                 respondsToSelector:@selector(amplitudeDidFinishInitializing:)]) {
+        [middleware amplitudeDidFinishInitializing:amplitude];
+    }
+}
+
+- (void)dispatchAmplitude:(Amplitude *)amplitude didUploadEventsManually:(BOOL)isManualUpload {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        if ([AMPMiddlewareRunner object:middleware
+                     respondsToSelector:@selector(amplitude:didUploadEventsManually:)]) {
+            [middleware amplitude:amplitude didUploadEventsManually:isManualUpload];
+        }
+    }
+}
+
+// AMPMiddleware never conformed to NSObject, which means we can't use the standard
+// [object respondsToSelector:] syntax to check for protocol conformance to optional methods.
++ (BOOL)object:(id)object respondsToSelector:(SEL)selector {
+    Class middlewareClass = object_getClass(object);
+    if (middlewareClass) {
+        return class_respondsToSelector(middlewareClass, selector);
+    }
+    return NO;
+}
+
 @end
diff --git a/Sources/Amplitude/Amplitude.m b/Sources/Amplitude/Amplitude.m
@@ -553,6 +553,8 @@ - (void)initializeApiKey:(NSString *)apiKey
                 [UIViewController amp_swizzleViewDidAppear];
             }
 #endif
+
+            [self->_middlewareRunner dispatchAmplitudeInitialized:self];
         }];
 
         if (!self.deferCheckInForeground) {
@@ -763,7 +765,7 @@ - (void)logEvent:(NSString *)eventType withEventProperties:(NSDictionary *)event
 
         int eventCount = [self.dbHelper getTotalEventCount]; // refetch since events may have been deleted
         if ((eventCount % self.eventUploadThreshold) == 0 && eventCount >= self.eventUploadThreshold) {
-            [self uploadEvents];
+            [self uploadEvents:NO];
         } else {
             [self uploadEventsWithDelay:self.eventUploadPeriodSeconds];
         }
@@ -956,20 +958,26 @@ - (void)uploadEventsWithDelay:(int)delay {
 
 - (void)uploadEventsInBackground {
     _updateScheduled = NO;
-    [self uploadEvents];
+    [self uploadEvents:NO];
 }
 
 - (void)uploadEvents {
+    [self uploadEvents:YES];
+}
+
+- (void)uploadEvents:(BOOL)isManualUpload {
     int limit = _backoffUpload ? _backoffUploadBatchSize : self.eventUploadMaxBatchSize;
-    [self uploadEventsWithLimit:limit];
+    [self uploadEventsWithLimit:limit isManualUpload:isManualUpload];
 }
 
-- (void)uploadEventsWithLimit:(int)limit {
+- (void)uploadEventsWithLimit:(int)limit isManualUpload:(BOOL)isManualUpload {
     if (self.apiKey == nil) {
         AMPLITUDE_ERROR(@"ERROR: apiKey cannot be nil or empty, set apiKey with initializeApiKey: before calling uploadEvents:");
         return;
     }
 
+    [_middlewareRunner dispatchAmplitude:self didUploadEventsManually:isManualUpload];
+
     @synchronized (self) {
         if (_updatingCurrently) {
             return;
@@ -1169,7 +1177,7 @@ - (void)makeEventUploadPostRequest:(NSString *)url events:(NSString *)events num
                 self->_backoffUploadBatchSize = MAX((int)ceilf(newNumEvents / 2.0f), 1);
                 AMPLITUDE_LOG(@"Request too large, will decrease size and attempt to reupload");
                 self->_updatingCurrently = NO;
-                [self uploadEventsWithLimit:self->_backoffUploadBatchSize];
+                [self uploadEventsWithLimit:self->_backoffUploadBatchSize isManualUpload:NO];
             } else if ([httpResponse statusCode] == 429) {
                 // if rate limited
                 self->_backoffUpload = YES;
@@ -1204,7 +1212,7 @@ - (void)makeEventUploadPostRequest:(NSString *)url events:(NSString *)events num
 
         if (uploadSuccessful && [self.dbHelper getEventCount] > self.eventUploadThreshold) {
             int limit = self->_backoffUpload ? self->_backoffUploadBatchSize : 0;
-            [self uploadEventsWithLimit:limit];
+            [self uploadEventsWithLimit:limit isManualUpload:NO];
     #if !TARGET_OS_OSX && !TARGET_OS_WATCH
         } else if (self->_uploadTaskID != UIBackgroundTaskInvalid) {
             if (uploadSuccessful) {
@@ -1243,7 +1251,7 @@ - (void)enterForeground {
 
         [self refreshDynamicConfig];
         [self startOrContinueSessionNSNumber:now inForeground:NO];
-        [self uploadEvents];
+        [self uploadEvents:NO];
     }];
 }
 
@@ -1267,7 +1275,7 @@ - (void)enterBackground {
 
     [self runOnBackgroundQueue:^{
         [self refreshSessionTime:now];
-        [self uploadEventsWithLimit:0];
+        [self uploadEventsWithLimit:0 isManualUpload:NO];
     }];
 }
 
@@ -1563,7 +1571,7 @@ - (void)setOffline:(BOOL)offline {
     _offline = offline;
 
     if (!_offline) {
-        [self uploadEvents];
+        [self uploadEvents:NO];
     }
 }
 
@@ -1646,6 +1654,9 @@ - (void)setServerZone:(AMPServerZone)serverZone updateServerUrl:(BOOL)updateServ
 
 - (void)addEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware {
     [_middlewareRunner add:middleware];
+    if (_initialized) {
+        [_middlewareRunner dispatchAmplitudeInitialized:self toMiddleware:middleware];
+    }
 }
 
 /**

diff --git a/Sources/Amplitude/Public/AMPMiddleware.h b/Sources/Amplitude/Public/AMPMiddleware.h
@@ -23,6 +23,8 @@
 
 #import <Foundation/Foundation.h>
 
+@class Amplitude;
+
 /**
  * AMPMiddlewarePayload
  */
@@ -44,6 +46,11 @@ typedef void (^AMPMiddlewareNext)(AMPMiddlewarePayload *_Nullable newPayload);
 
 - (void)run:(AMPMiddlewarePayload *_Nonnull)payload next:(AMPMiddlewareNext _Nonnull)next;
 
+@optional
+
+- (void)amplitudeDidFinishInitializing:(nonnull Amplitude *)amplitude;
+- (void)amplitude:(nonnull Amplitude *)amplitude didUploadEventsManually:(BOOL)manually;
+
 @end
 
 /**
@@ -55,6 +62,9 @@ typedef void (^AMPMiddlewareBlock)(AMPMiddlewarePayload *_Nonnull payload, AMPMi
 
 @property (nonnull, nonatomic, readonly) AMPMiddlewareBlock block;
 
-- (instancetype _Nonnull)initWithBlock:(AMPMiddlewareBlock _Nonnull)block;
+@property (nonatomic, copy, nullable) void (^didFinishInitializing)(Amplitude * _Nonnull amplitude);
+@property (nonatomic, copy, nullable) void (^didUploadEventsManually)(Amplitude * _Nonnull amplitude, BOOL isManualUpload);
+
+- (instancetype _Nonnull)initWithBlock:(AMPMiddlewareBlock _Nonnull)block NS_DESIGNATED_INITIALIZER;
 
 @end
diff --git a/Tests/Amplitude+Test.h b/Tests/Amplitude+Test.h
@@ -37,5 +37,6 @@
 - (NSDate*)currentTime;
 - (id)unarchive:(NSString*)path;
 - (BOOL)archive:(id)obj toFile:(NSString*)path;
+- (void)removeObservers;
 
 @end
diff --git a/Tests/AmplitudeTests.m b/Tests/AmplitudeTests.m
@@ -1246,6 +1246,82 @@ - (void)testMiddlewareSupport {
     XCTAssertEqualObjects(middlewareExtra[@"description"], event[@"description"]);
 }
 
+- (void)testMiddlewareInitializeEvents {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_lifecyle_support"];
+
+    const XCTestExpectation *preInitMiddlewareExpectation = [self expectationWithDescription:@"calls intialized"];
+    const AMPBlockMiddleware *preInitMiddleware = [[AMPBlockMiddleware alloc] init];
+    preInitMiddleware.didFinishInitializing = ^(Amplitude *amplitude) {
+        [preInitMiddlewareExpectation fulfill];
+    };
+    [client addEventMiddleware:preInitMiddleware];
+
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    const XCTestExpectation *postInitMiddlewareExpectation = [self expectationWithDescription:@"calls intialized"];
+    const AMPBlockMiddleware *postInitMiddleware = [[AMPBlockMiddleware alloc] init];
+    postInitMiddleware.didFinishInitializing = ^(Amplitude *amplitude) {
+        [postInitMiddlewareExpectation fulfill];
+    };
+    [client addEventMiddleware:postInitMiddleware];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testMiddlewareManualFlush {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_manual_flush"];
+    // prevent automatic flushes from parallel tests
+    [client removeObservers];
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    const XCTestExpectation *manualFlushMiddlewareExpectation = [self expectationWithDescription:@"calls flush"];
+    const AMPBlockMiddleware *manualFlushMiddleware = [[AMPBlockMiddleware alloc] init];
+    manualFlushMiddleware.didUploadEventsManually = ^(Amplitude * _Nonnull amplitude, BOOL isManualUpload) {
+        XCTAssertTrue(isManualUpload);
+        [manualFlushMiddlewareExpectation fulfill];
+    };
+    [client addEventMiddleware:manualFlushMiddleware];
+
+    [client uploadEvents];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testMiddlewareAutomaticFlush {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_automatic_flush"];
+    // prevent automatic flushes from parallel tests
+    [client removeObservers];
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    const XCTestExpectation *automaticFlushMiddlewareExpectation = [self expectationWithDescription:@"calls flush"];
+    const AMPBlockMiddleware *automaticFlushMiddleware = [[AMPBlockMiddleware alloc] init];
+    automaticFlushMiddleware.didUploadEventsManually = ^(Amplitude * _Nonnull amplitude, BOOL isManualUpload) {
+        XCTAssertFalse(isManualUpload);
+        [automaticFlushMiddlewareExpectation fulfill];
+    };
+    [client addEventMiddleware:automaticFlushMiddleware];
+
+    client.eventUploadThreshold = 1;
+    [client logEvent:@"test"];
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
 - (void)testSwallowMiddleware {
     AMPBlockMiddleware *swallowMiddleware = [[AMPBlockMiddleware alloc] initWithBlock: ^(AMPMiddlewarePayload * _Nonnull payload, AMPMiddlewareNext _Nonnull next) {
     }];

diff --git a/Tests/MiddlewareRunnerTests.m b/Tests/MiddlewareRunnerTests.m
@@ -7,6 +7,7 @@
 //
 
 #import <XCTest/XCTest.h>
+#import "Amplitude/Amplitude.h"
 #import "AMPMiddleware.h"
 #import "AMPMiddlewareRunner.h"
 
@@ -87,4 +88,54 @@ - (void)testRunWithNotPassMiddleware {
     XCTAssertEqualObjects([event objectForKey:@"event_type"], eventType);
 }
 
+- (void)testSendsAmplitudeInitialized {
+    const XCTestExpectation *didSendDidFinishInitializingExpectation = [self expectationWithDescription:@"didSendFinishInitializing"];
+    const AMPBlockMiddleware *middleware = [[AMPBlockMiddleware alloc] init];
+    middleware.didFinishInitializing = ^(Amplitude *amplitude) {
+        [didSendDidFinishInitializingExpectation fulfill];
+    };
+    [self.middlewareRunner add:middleware];
+    [self.middlewareRunner dispatchAmplitudeInitialized:[[Amplitude alloc] init]];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:^(NSError *error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testSendsAmplitudeInitializedToMiddleware {
+    const XCTestExpectation *didSendDidFinishInitializingExpectation = [self expectationWithDescription:@"didSendFinishInitializing"];
+    const AMPBlockMiddleware *middleware = [[AMPBlockMiddleware alloc] init];
+    middleware.didFinishInitializing = ^(Amplitude *amplitude) {
+        [didSendDidFinishInitializingExpectation fulfill];
+    };
+    [self.middlewareRunner add:middleware];
+    [self.middlewareRunner dispatchAmplitudeInitialized:[[Amplitude alloc] init] toMiddleware:middleware];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:^(NSError *error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testSendsAmplitudeDidSendFlush {
+    const BOOL manualUpload = YES;
+    const XCTestExpectation *didSendFlushExpectation = [self expectationWithDescription:@"didSendFlush"];
+    const AMPBlockMiddleware *middleware = [[AMPBlockMiddleware alloc] init];
+    middleware.didUploadEventsManually = ^(Amplitude *amplitude, BOOL isManualUpload) {
+        XCTAssertEqual(isManualUpload, manualUpload);
+        [didSendFlushExpectation fulfill];
+    };
+    [self.middlewareRunner add:middleware];
+    [self.middlewareRunner dispatchAmplitude:[[Amplitude alloc] init] didUploadEventsManually:manualUpload];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:^(NSError *error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
 @end