Keyed NSOperationQueue

KeyedOperationQueue.h

#import <Foundation/Foundation.h>


extern id KeyedOperationQueueCreateUniqueKey();

@interface KeyedOperationQueue : NSOperationQueue

@property (nonatomic, readonly) NSDictionary *operationsForKeys;

- (void)addOperation:(NSOperation *)operation forKey:(id)key;
- (void)addOperationWithBlock:(void (^)(void))block forKey:(id)key;

- (id)addOperationForUniqueKey:(NSOperation *)operation;
- (id)addOperationWithBlockForUniqueKey:(void (^)(void))block;

- (void)cancelOperationForKey:(id)key;

@end

KeyedOperationQueue.m

#import "KeyedOperationQueue.h"


id KeyedOperationQueueCreateUniqueKey()
{
    CFUUIDRef uuid = CFUUIDCreate(NULL);
    CFStringRef string = CFUUIDCreateString(NULL, uuid);
    CFRelease(uuid);
    return (__bridge_transfer NSString *)string;
}

@implementation KeyedOperationQueue {
    NSMutableDictionary *_operationsForKeys;
}

- (id)init
{
    self = [super init];
    if (self) {
        _operationsForKeys = [NSMutableDictionary new];
    }
    return self;
}

- (void)addOperation:(NSOperation *)operation forKey:(id)key
{
    if (key) {
        [self cancelOperationForKey:key];
        
        __weak NSOperation *weakOperation = operation;
        void(^completionBlock)(void) = ^{
            @synchronized(_operationsForKeys) {
                NSSet *keys = [_operationsForKeys keysOfEntriesWithOptions:NSEnumerationReverse passingTest:^BOOL(id key, id obj, BOOL *stop) {
                    return obj == weakOperation;
                }];
                [_operationsForKeys removeObjectsForKeys:[keys allObjects]];
            }
        };
        
        void(^originalBlock)(void) = [operation completionBlock];
        void(^replacedBlock)(void);
        
        if (originalBlock) {
            replacedBlock = ^{
                originalBlock();
                completionBlock();
            };
        } else {
            replacedBlock = completionBlock;
        }
        
        [operation setCompletionBlock:replacedBlock];
        
        @synchronized(_operationsForKeys) {
            [_operationsForKeys setObject:operation forKey:key];
        }
    }
    
    [self addOperation:operation];
}

- (void)addOperationWithBlock:(void (^)(void))block forKey:(id)key
{
    NSOperation *operation = [NSBlockOperation blockOperationWithBlock:block];
    [self addOperation:operation forKey:key];
}

- (id)addOperationForUniqueKey:(NSOperation *)operation
{
    id key = KeyedOperationQueueCreateUniqueKey();
    [self addOperation:operation forKey:key];
    return key;
}

- (id)addOperationWithBlockForUniqueKey:(void (^)(void))block
{
    id key = KeyedOperationQueueCreateUniqueKey();
    [self addOperationWithBlock:block forKey:key];
    return key;
}

- (void)cancelOperationForKey:(id)key
{
    if (!key) {
        return;
    }
    
    @synchronized(_operationsForKeys) {
        NSOperation *operation = [_operationsForKeys objectForKey:key];
        [operation cancel];
        [_operationsForKeys removeObjectForKey:key];
    }
}

- (void)cancelAllOperations
{
    [super cancelAllOperations];
    
    @synchronized(_operationsForKeys) {
        [_operationsForKeys removeAllObjects];
    }
}

@end