PINOperationGroup.m 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. //
  2. // PINOperationGroup.m
  3. // PINQueue
  4. //
  5. // Created by Garrett Moon on 10/8/16.
  6. // Copyright © 2016 Pinterest. All rights reserved.
  7. //
  8. #import "PINOperationGroup.h"
  9. #import "PINOperation.h"
  10. #import <pthread.h>
  11. @interface NSNumber (PINGroupOperationQueue) <PINGroupOperationReference>
  12. @end
  13. @interface PINOperationGroup ()
  14. {
  15. pthread_mutex_t _lock;
  16. PINOperationQueue *_operationQueue;
  17. NSMutableArray <dispatch_block_t> *_operations;
  18. NSMutableArray <NSNumber *> *_operationPriorities;
  19. NSMutableArray <id <PINGroupOperationReference>> *_operationReferences;
  20. NSMapTable <id <PINGroupOperationReference>, id <PINOperationReference>> *_groupToOperationReferences;
  21. NSUInteger _operationReferenceCount;
  22. dispatch_group_t _group;
  23. dispatch_block_t _completion;
  24. dispatch_queue_t _completionQueue;
  25. BOOL _started;
  26. BOOL _canceled;
  27. }
  28. - (instancetype)initWithOperationQueue:(PINOperationQueue *)operationQueue NS_DESIGNATED_INITIALIZER;
  29. @end
  30. @implementation PINOperationGroup
  31. - (instancetype)initWithOperationQueue:(PINOperationQueue *)operationQueue
  32. {
  33. if (self = [super init]) {
  34. pthread_mutex_init(&_lock, NULL);
  35. _operationQueue = operationQueue;
  36. _operations = [[NSMutableArray alloc] init];
  37. _operationReferences = [[NSMutableArray alloc] init];
  38. _operationPriorities = [[NSMutableArray alloc] init];
  39. _groupToOperationReferences = [NSMapTable weakToStrongObjectsMapTable];
  40. _group = dispatch_group_create();
  41. }
  42. return self;
  43. }
  44. - (void)dealloc
  45. {
  46. pthread_mutex_destroy(&_lock);
  47. }
  48. + (instancetype)asyncOperationGroupWithQueue:(PINOperationQueue *)operationQueue
  49. {
  50. return [[self alloc] initWithOperationQueue:operationQueue];
  51. }
  52. - (id <PINGroupOperationReference>)locked_nextOperationReference
  53. {
  54. id <PINGroupOperationReference> reference = [NSNumber numberWithUnsignedInteger:++_operationReferenceCount];
  55. return reference;
  56. }
  57. - (void)start
  58. {
  59. [self lock];
  60. NSAssert(_canceled == NO, @"Operation group canceled.");
  61. if (_started == NO && _canceled == NO) {
  62. for (NSUInteger idx = 0; idx < _operations.count; idx++) {
  63. dispatch_group_enter(_group);
  64. dispatch_block_t originalOperation = _operations[idx];
  65. dispatch_block_t groupBlock = ^{
  66. originalOperation();
  67. dispatch_group_leave(_group);
  68. };
  69. id <PINOperationReference> operationReference = [_operationQueue addOperation:groupBlock withPriority:[_operationPriorities[idx] unsignedIntegerValue]];
  70. [_groupToOperationReferences setObject:operationReference forKey:_operationReferences[idx]];
  71. }
  72. if (_completion) {
  73. dispatch_queue_t completionQueue = _completionQueue ? _completionQueue : dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
  74. dispatch_group_notify(_group, completionQueue, ^{
  75. dispatch_block_t completion;
  76. [self lock];
  77. completion = _completion;
  78. _completion = nil;
  79. [self unlock];
  80. if (completion) {
  81. completion();
  82. }
  83. });
  84. }
  85. _operations = nil;
  86. _operationPriorities = nil;
  87. _operationReferences = nil;
  88. }
  89. [self unlock];
  90. }
  91. - (void)cancel
  92. {
  93. [self lock];
  94. _canceled = YES;
  95. for (id <PINOperationReference>operationReference in [_groupToOperationReferences objectEnumerator]) {
  96. if ([_operationQueue cancelOperation:operationReference]) {
  97. dispatch_group_leave(_group);
  98. }
  99. }
  100. //TODO just nil out instead? Does it make sense to support adding operations after cancelation?
  101. [_groupToOperationReferences removeAllObjects];
  102. [_operations removeAllObjects];
  103. [_operationPriorities removeAllObjects];
  104. [_operationReferences removeAllObjects];
  105. _completion = nil;
  106. [self unlock];
  107. }
  108. - (id <PINGroupOperationReference>)addOperation:(dispatch_block_t)operation
  109. {
  110. return [self addOperation:operation withPriority:PINOperationQueuePriorityDefault];
  111. }
  112. - (id <PINGroupOperationReference>)addOperation:(dispatch_block_t)operation withPriority:(PINOperationQueuePriority)priority
  113. {
  114. [self lock];
  115. id <PINGroupOperationReference> reference = nil;
  116. NSAssert(_started == NO && _canceled == NO, @"Operation group already started or canceled.");
  117. if (_started == NO && _canceled == NO) {
  118. reference = [self locked_nextOperationReference];
  119. [_operations addObject:operation];
  120. [_operationPriorities addObject:@(priority)];
  121. [_operationReferences addObject:reference];
  122. }
  123. [self unlock];
  124. return reference;
  125. }
  126. - (void)setCompletion:(dispatch_block_t)completion
  127. {
  128. [self lock];
  129. NSAssert(_started == NO && _canceled == NO, @"Operation group already started or canceled.");
  130. if (_started == NO && _canceled == NO) {
  131. _completion = completion;
  132. }
  133. [self unlock];
  134. }
  135. - (void)waitUntilComplete
  136. {
  137. [self start];
  138. dispatch_group_wait(_group, DISPATCH_TIME_FOREVER);
  139. }
  140. - (void)lock
  141. {
  142. pthread_mutex_lock(&_lock);
  143. }
  144. - (void)unlock
  145. {
  146. pthread_mutex_unlock(&_lock);
  147. }
  148. @end