DMMemoize/DMMemoize.m at master · delicious-monster/DMMemoize · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
//
//  DMMemoize.m
//  DMMemoize
//
//  Created by Jonathon Mah on 2012-11-10.
//  Copyright (c) 2012 Delicious Monster Software. All rights reserved.
//

#import "DMMemoize.h"

#import <libkern/OSAtomic.h>
#import <objc/runtime.h>


@implementation DMMemoize

+ (id)cachedValueForKey:(id)cacheKey storageOwner:(id)cacheOwner token:(void *)staticToken generator:(id(^)(void))generatorBlock;
{
    static dispatch_once_t onceToken;
    static id nilMarker;
    dispatch_once(&onceToken, ^{
        nilMarker = [NSObject new];
    });

    NSParameterAssert(cacheOwner && staticToken && generatorBlock);
    if (!cacheOwner || !staticToken)
        return nil;

    const id nonNilCacheKey = cacheKey ? : nilMarker;

    static OSSpinLock globalSpinLock = OS_SPINLOCK_INIT; // locked while looking up or creating the per-owner mutex
    // Take the global lock during creation or look-up of finer-grained cacheStorage mutex
    OSSpinLockLock(&globalSpinLock);
    NSMapTable *cacheStorage = objc_getAssociatedObject(cacheOwner, staticToken);
    if (!cacheStorage) {
        cacheStorage = [NSMapTable strongToStrongObjectsMapTable];
        objc_setAssociatedObject(cacheOwner, staticToken, cacheStorage, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
    }

    static char storageMutexAssociationKey;
    dispatch_semaphore_t cacheStorageMutex = objc_getAssociatedObject(cacheStorage, &storageMutexAssociationKey);
    if (!cacheStorageMutex) {
        cacheStorageMutex = dispatch_semaphore_create(1);
        objc_setAssociatedObject(cacheStorage, &storageMutexAssociationKey, cacheStorageMutex, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
    }
    OSSpinLockUnlock(&globalSpinLock);


    dispatch_semaphore_wait(cacheStorageMutex, DISPATCH_TIME_FOREVER);

    id cachedValue = [cacheStorage objectForKey:nonNilCacheKey];
    if (!cachedValue) {
        // Value isn't cached. Another thread may already be generating it though.
        static char inProgressSemaphoreTableAssociationKey;
        NSMapTable *generatorInProgressSemaphores = objc_getAssociatedObject(cacheStorage, &inProgressSemaphoreTableAssociationKey);
        if (!generatorInProgressSemaphores) {
            generatorInProgressSemaphores = [NSMapTable strongToStrongObjectsMapTable];
            objc_setAssociatedObject(cacheStorage, &inProgressSemaphoreTableAssociationKey, generatorInProgressSemaphores, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
        }

        NSMutableArray *waitingSemaphoresArray = [generatorInProgressSemaphores objectForKey:nonNilCacheKey];
        if (waitingSemaphoresArray) {
            // Another thread is generating the value for this key; wait until it's done.
            dispatch_semaphore_t myWaitingSemaphore = dispatch_semaphore_create(0);
            [waitingSemaphoresArray addObject:myWaitingSemaphore];
            dispatch_semaphore_signal(cacheStorageMutex); // unlock while we wait
            dispatch_semaphore_wait(myWaitingSemaphore, DISPATCH_TIME_FOREVER);
            dispatch_semaphore_wait(cacheStorageMutex, DISPATCH_TIME_FOREVER);

            cachedValue = [cacheStorage objectForKey:nonNilCacheKey];

        } else {
            // We are responsible for generating this key. Create a semaphore so others can wait on us.
            waitingSemaphoresArray = [NSMutableArray array];
            [generatorInProgressSemaphores setObject:waitingSemaphoresArray forKey:nonNilCacheKey];

            // Unlock before we call back into client code, because the generator could use other cached values
            dispatch_semaphore_signal(cacheStorageMutex);
            cachedValue = generatorBlock() ? : nilMarker;
            dispatch_semaphore_wait(cacheStorageMutex, DISPATCH_TIME_FOREVER);

            [cacheStorage setObject:cachedValue forKey:nonNilCacheKey];

            [generatorInProgressSemaphores removeObjectForKey:nonNilCacheKey];
            if ([generatorInProgressSemaphores count] == 0)
                // Dispose of empty semaphore table
                objc_setAssociatedObject(cacheStorage, &inProgressSemaphoreTableAssociationKey, nil, OBJC_ASSOCIATION_RETAIN_NONATOMIC);

            for (dispatch_semaphore_t waitingSemaphore in waitingSemaphoresArray)
                dispatch_semaphore_signal(waitingSemaphore); // signal all waiters (though they will immediately go back to sleep on cacheStorageMutex)
        }
    }
    dispatch_semaphore_signal(cacheStorageMutex);

    return (cachedValue != nilMarker) ? cachedValue : nil;
}

@end