Automatic removal of NSNotificationCenter or KVO observers

Observer pattern is common on iOS platform, you use observers in NSNotificationCenter and in Key Value Observing.

You need to remember to unregister before you release your observer object, if you don’t you are going to have crashes.

What if you could automate it ?

Idea

There are actually 2 categories, one for NSNotificationCenter and second for KVO. They use very similar code so I will just explain a few tricks used to make it work. You are free to look at whole code but I will explain only key points…

Since I wanted to have a category that doesn’t require changing code in existing apps, that meant that the methods you are going to use must stay the same, no custom methods calling or additional code needed to be called. That means I have to swizzle the existing methods with my own.

It should happen only once and it should happen automaticlly, my first thought was

+(void)initialize

this method is called when a class is used for first time, but I’m not subclassing but creating categories so it wouldn’t work correctly. Instead I decided to use -(void)load as it gets called even for categories. But remember that load is called before main, so you don’t have normal autorelease pool, you need to create one:

- (void)load
{
  //! swap methods
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    @autoreleasepool {
      [self sf_swapSelector:@selector(addObserver:forKeyPath:options:context:) withSelector:@selector(sf_addObserver:forKeyPath:options:context:)];
      [self sf_swapSelector:@selector(removeObserver:forKeyPath:) withSelector:@selector(sf_removeObserver:forKeyPath:)];
      [self sf_swapSelector:@selector(removeObserver:forKeyPath:context:) withSelector:@selector(sf_removeObserver:forKeyPath:context:)];
    }
  });
}

sf_swapSelector changes default selectors to call sf_ prefixed ones, it also creates a new selectors called sf_original_”originalSelector” that point to original code.

Data storage

Next time someone adds / removes observer we need to store that data so that we can operate on it later. Since we need this data when working with the observer object we should store it inside it. We can use associated references for that.

I have used Dictionary of Arrays ( key is actually keyPath ) of internal object named __SFObserverKVOObserverInfo that looks as simple as that:

@interface __SFObserversKVOObserverInfo : NSObject
@property(nonatomic, copy) NSString *keyPath;
@property(nonatomic, AH_WEAK) id context;
@property(nonatomic, assign) void *blockKey;
@end

keyPath and context are self explanatory, blockKey is a key used to access block created by SFExecuteOnDealloc( very handy category that allows you to execute any block when an object get’s deallocated, useful especially when you create advanced categories like this one ).

Execute on dealloc

Each time we add new observer, we need to make sure that it will be deallocated automatically, this is the place that SFExecuteOnDealloc get’s really useful:

//! 1
__unsafe_unretained __block id weakSelf = self;
__unsafe_unretained __block id weakObserver = observer;
__unsafe_unretained __block id weakContext = aContext;
//! 2
void *key = [observer performBlockOnDealloc:^{
//! 3
  if ([weakSelf sf_removeObserver:weakObserver forKeyPath:keyPath context:weakContext registeredKeyPaths:registeredKeyPaths]) {
    //! 4
    [self setAllowMethodForwarding:YES];
    //! 5
    objc_msgSend(self, NSSelectorFromString(NSObjectKVOSFObserversRemoveSpecificSelector), weakObserver, keyPath, weakContext);
    [self setAllowMethodForwarding:NO];
  }
}];

1. **__unsafe_unretained __block **specifier breaks retain cycles on both ARC and non ARC code, and since we are using block we need to make sure we do not create retain cycle ( any obj-c object without it inside a block will automatically be retained as long as the block exists )
2. We add a new block to be executed when observer get’s deallocated
3. We check and remove data from our internal structure for specified options, it will return YES if something was actually removed ( prevents removal of objects that isn’t observer anymore )
4. Enables method forwarding, more on this later…
5. Calls original method, we are using** NSSelectorFromString to prevent compiler warnings about unknown selectors, just a clever trick I came up with.**

Method Forwarding

Let me explain what happens in point 4.

When we call our original selector, theoretically we skip our swizzled function and go straight to original implementation,** BUT the system calls from this original method other methods we swizzled **for example when you just call removeObserver: for NSNotificationCenter it may call more specific removeObserver:name:object: for each object that is registered.

We swizzled original selector to point to our method, so now its important that we proceed carefully. There are 2 cases our removeObserver: methods can be called:

1. Programmers call removeObserver somewhere from the application code.
2. The methods get called due to original implementation calling different methods. (This will be identified as our method forwarding).

We could try to mimic original behavior in our code by analyzing every call and recreating the algorithm they use but this could break if apple changes something, instead I figured out that since code is executed synchronously we can identify both cases, that’s why we have allowMethodForwarding property. Look at the code of our removeObserver:

- (void)sf_removeObserver:(id)observer forKeyPath:(NSString *)keyPath
{
//! 1
  if ([self allowMethodForwarding]) {
#if SF_OBSERVERS_LOG_ORIGINAL_METHODS
    NSLog(@"Calling original method %@ with parameters %@ %@", NSObjectKVOSFObserversRemoveSelector, observer, keyPath);
#endif
    objc_msgSend(self, NSSelectorFromString(NSObjectKVOSFObserversRemoveSelector), observer, keyPath);
    return;
  }
//! 2
  NSMutableDictionary *registeredKeyPaths = (NSMutableDictionary *)objc_getAssociatedObject(observer, AH_BRIDGE(NSObjectKVOSFObserversArrayKey));
  if ([self sf_removeObserver:observer forKeyPath:keyPath context:nil registeredKeyPaths:registeredKeyPaths]) {
#if SF_OBSERVERS_LOG_ORIGINAL_METHODS
      NSLog(@"Calling original method %@ with parameters %@ %@", NSObjectKVOSFObserversRemoveSelector, observer, keyPath);
#endif
//! 3
    [self setAllowMethodForwarding:YES];
    objc_msgSend(self, NSSelectorFromString(NSObjectKVOSFObserversRemoveSelector), observer, keyPath);
    [self setAllowMethodForwarding:NO];
  }
}

1. If we are method forwarding just call original selector and return.
2. Check our internal structure if we can remove object ( we prevent removing object that isn’t observer ).
3. Since this is called by our code and not by method forwarding, enable method forwarding in case original code calls other swizzled selectors then disable it afterwards.

Conclusion

You no longer need to call removeObserver in your dealloc methods, that especially useful when using ARC as you won’t need to create dealloc at all in 95 % classes in your application. You also learned, that you could use NSSelectorFromString to prevent compiler from generating warning for dynamically created methods, that __unsafe_unretained __block specifier is the correct way to break retain cycles on code that needs to be run both on ARC and non ARC.

Also you can have seen how to be on the safe side when swizzling multiples methods that can call each other by identifing methodForwarding, instead of trying to recreate system behavior that could change at any iOS update.

grab it from GitHub