Skip to content

Commit

Permalink
feat: extend middleware for session replay (#501)
Browse files Browse the repository at this point in the history
  • Loading branch information
crleona authored Aug 12, 2024
1 parent b7a66d2 commit 051890f
Show file tree
Hide file tree
Showing 9 changed files with 226 additions and 10 deletions.
18 changes: 18 additions & 0 deletions Sources/Amplitude/AMPMiddleware.m
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
9 changes: 9 additions & 0 deletions Sources/Amplitude/AMPMiddlewareRunner.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
#import <Foundation/Foundation.h>
#import "AMPMiddleware.h"

@class Amplitude;

@interface AMPMiddlewareRunner : NSObject

@property (nonatomic, nonnull, readonly) NSMutableArray<id<AMPMiddleware>> *middlewares;
Expand All @@ -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
34 changes: 34 additions & 0 deletions Sources/Amplitude/AMPMiddlewareRunner.m
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
//

#import <Foundation/Foundation.h>
#import <objc/runtime.h>
#import "AMPMiddlewareRunner.h"
#import "AMPMiddleware.h"

Expand Down Expand Up @@ -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
33 changes: 24 additions & 9 deletions Sources/Amplitude/Amplitude.m
Original file line number Diff line number Diff line change
Expand Up @@ -553,6 +553,8 @@ - (void)initializeApiKey:(NSString *)apiKey
[UIViewController amp_swizzleViewDidAppear];
}
#endif

[self->_middlewareRunner dispatchAmplitudeInitialized:self];
}];

if (!self.deferCheckInForeground) {
Expand Down Expand Up @@ -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];
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -1243,7 +1251,7 @@ - (void)enterForeground {

[self refreshDynamicConfig];
[self startOrContinueSessionNSNumber:now inForeground:NO];
[self uploadEvents];
[self uploadEvents:NO];
}];
}

Expand All @@ -1267,7 +1275,7 @@ - (void)enterBackground {

[self runOnBackgroundQueue:^{
[self refreshSessionTime:now];
[self uploadEventsWithLimit:0];
[self uploadEventsWithLimit:0 isManualUpload:NO];
}];
}

Expand Down Expand Up @@ -1563,7 +1571,7 @@ - (void)setOffline:(BOOL)offline {
_offline = offline;

if (!_offline) {
[self uploadEvents];
[self uploadEvents:NO];
}
}

Expand Down Expand Up @@ -1633,6 +1641,10 @@ - (void)setIngestionMetadata:(AMPIngestionMetadata *)ingestionMetadata {
_ingestionMetadata = ingestionMetadata;
}

- (AMPServerZone)serverZone {
return _serverZone;
}

- (void)setServerZone:(AMPServerZone)serverZone {
[self setServerZone:serverZone updateServerUrl:YES];
}
Expand All @@ -1646,6 +1658,9 @@ - (void)setServerZone:(AMPServerZone)serverZone updateServerUrl:(BOOL)updateServ

- (void)addEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware {
[_middlewareRunner add:middleware];
if (_initialized) {
[_middlewareRunner dispatchAmplitudeInitialized:self toMiddleware:middleware];
}
}

/**
Expand Down
12 changes: 11 additions & 1 deletion Sources/Amplitude/Public/AMPMiddleware.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@

#import <Foundation/Foundation.h>

@class Amplitude;

/**
* AMPMiddlewarePayload
*/
Expand All @@ -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

/**
Expand All @@ -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
2 changes: 2 additions & 0 deletions Sources/Amplitude/Public/Amplitude.h
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,8 @@ typedef void (^AMPInitCompletionBlock)(void);

- (void)setIngestionMetadata:(AMPIngestionMetadata *)ingestionMetadata;

- (AMPServerZone)serverZone;

/**
* Set Amplitude Server Zone, switch to zone related configuration, including dynamic configuration and server url.
* To send data to Amplitude's EU servers, you need to configure the serverZone to EU like [client setServerZone:EU]
Expand Down
1 change: 1 addition & 0 deletions Tests/Amplitude+Test.h
Original file line number Diff line number Diff line change
Expand Up @@ -37,5 +37,6 @@
- (NSDate*)currentTime;
- (id)unarchive:(NSString*)path;
- (BOOL)archive:(id)obj toFile:(NSString*)path;
- (void)removeObservers;

@end
76 changes: 76 additions & 0 deletions Tests/AmplitudeTests.m
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
}];
Expand Down
Loading

0 comments on commit 051890f

Please sign in to comment.