fixed cache leak when context contains non-recreatable/invalid context

Summary: This change allows reuse pool to "ask" controller if context is still valid and can be recreated for future reuse. Goal is to fix cache leak issue, when view is stored in reuse pool but would be never reused, since context contains piece of data that can't be recreated anymore.

Reviewed By: kevin0571

Differential Revision: D26914733

fbshipit-source-id: de4247b78328631b9407a9f5a844a3d535e4972c

Author: mkareta

ComponentKit/StatefulViews/CKStatefulViewComponentController.h

@@ -61,6 +61,17 @@
  */
 + (NSInteger)maximumPoolSize:(ContextType)context;
 
+/**
+ Optionally override this to validate if context is still valid and related view can be reused.
+ If context becames invalid, it will be eventually removed from reuse pool.
+ Default implementaion returns YES.
+ This method should return NO, only when context can't be recreated anymore,
+ and view will be never reused.
+ WARNING: context invalidation must be one way process. If this function
+ returns YES for context that was previously invalid then behaviour is undefined.
+ */
++ (BOOL)isContextValid:(ContextType)context;
+
 /**
  The current stateful view owned by this controller, if any.
 

ComponentKit/StatefulViews/CKStatefulViewComponentController.mm

@@ -47,6 +47,10 @@ + (NSInteger)maximumPoolSize:(id)context
   return -1;
 }
 
++ (BOOL)isContextValid:(id)context {
+  return YES;
+}
+
 - (UIView *)statefulView
 {
   return _statefulView;

ComponentKit/StatefulViews/CKStatefulViewReusePool.mm

@@ -123,6 +123,16 @@ - (UIView *)dequeueStatefulViewForControllerClass:(Class)controllerClass
   if (it == _pool.end()) { // Avoid overhead of creating the item unless it already exists
     return nil;
   }
+  if (![it->first.first isContextValid:it->first.second]) {
+    // Context is invalid. This generally should not happen
+    // since invalid context == non reacreatable context
+    // But we still handle it to make sure that behaviour is
+    // well defined
+
+    // Remove views with invalid context from reuse pool
+    _pool.erase(it);
+    return nil;
+  }
   UIView *candidate = it->second.viewWithPreferredSuperview(preferredSuperview);
   if (candidate) {
     return candidate;
@@ -144,6 +154,13 @@ - (void)enqueueStatefulView:(UIView *)view
   RCAssertNotNil(controllerClass, @"Must provide a controller class");
   RCAssertNotNil(mayRelinquishBlock, @"Must provide a relinquish block");
 
+  // Shortcut for invalid contexts. If context is invalid
+  // by the time when it is being added to the pool, we should
+  // drop it immediately
+  if (![controllerClass isContextValid:context]) {
+    return;
+  }
+
   auto const addEntry = ^{
     auto &poolItem = _pendingPool[std::make_pair(controllerClass, context)];
     poolItem.addEntry({view, mayRelinquishBlock});
@@ -165,13 +182,20 @@ - (void)enqueueStatefulView:(UIView *)view
   RCDispatchMainDefaultMode(^{
     self->_enqueuedPendingPurge = NO;
     [self purgePendingPool];
+    [self dropViewsWithInvalidContext];
   });
 }
 
 - (void)purgePendingPool
 {
   RCAssertMainThread();
   for (const auto &it : _pendingPool) {
+    // Ignore items that already can't be reused to save some cycles
+    BOOL isContextValid = [it.first.first isContextValid:it.first.second];
+    if (!isContextValid) {
+      continue;
+    }
+
     // maximumPoolSize will be -1 by default
     NSInteger maximumPoolSize = [it.first.first maximumPoolSize:it.first.second];
 
@@ -183,4 +207,15 @@ - (void)purgePendingPool
   _clearingPendingPool = NO;
 }
 
+- (void)dropViewsWithInvalidContext {
+  RCAssertMainThread();
+  for (auto it = _pool.begin(); it != _pool.end();) {
+    if (![it->first.first isContextValid:it->first.second]) {
+      it = _pool.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 @end

ComponentKitTests/StatefulViews/CKStatefulViewReusePoolTests.mm

@@ -34,6 +34,22 @@ + (NSInteger)maximumPoolSize:(id)context
 
 @end
 
+@interface CKStatefulViewComponentInvalidatableContext: NSObject
+@property (assign, nonatomic) BOOL isValid;
+@end
+@implementation CKStatefulViewComponentInvalidatableContext
+@end
+
+@interface CKStatefulViewComponentWithInvalidContextController : CKStatefulViewComponentController
+@end
+@implementation CKStatefulViewComponentWithInvalidContextController
+
++ (BOOL)isContextValid:(CKStatefulViewComponentInvalidatableContext *)context {
+  return context.isValid;
+}
+
+@end
+
 @interface EnqueueOnDealloc: NSObject
 + (instancetype)newWithPool:(CKStatefulViewReusePool *)pool;
 @end
@@ -272,6 +288,95 @@ - (void)testMaximumPoolSizeOfOneByEnqueueingTwoViewsThenDequeueingTwoViewsReturn
   XCTAssertTrue(dequeuedView2 != view2, @"Didn't expect view in container2 to be returned");
 }
 
+- (void)testInvalidContextEnqueue
+{
+  CKStatefulViewReusePool *pool = [[CKStatefulViewReusePool alloc] init];
+  CKStatefulViewComponentInvalidatableContext *context = [CKStatefulViewComponentInvalidatableContext new];
+
+  __block int calledBlockCount = 0;
+
+  UIView *container1 = [[UIView alloc] init];
+  CKTestStatefulView *view1 = [[CKTestStatefulView alloc] init];
+  [container1 addSubview:view1];
+  [pool enqueueStatefulView:view1
+         forControllerClass:[CKStatefulViewComponentWithInvalidContextController class]
+                    context:context
+         mayRelinquishBlock:^BOOL{
+           return YES;
+         }];
+
+  UIView *container2 = [[UIView alloc] init];
+  CKTestStatefulView *view2 = [[CKTestStatefulView alloc] init];
+  [container2 addSubview:view2];
+  [pool enqueueStatefulView:view2
+         forControllerClass:[CKOtherStatefulViewComponentController class]
+                    context:nil
+         mayRelinquishBlock:^BOOL{
+           calledBlockCount++;
+           return YES;
+         }];
+
+  CKRunRunLoopUntilBlockIsTrue(^BOOL{
+    return calledBlockCount == 1;
+  });
+
+  UIView *dequeuedView1 = [pool dequeueStatefulViewForControllerClass:[CKStatefulViewComponentWithInvalidContextController class]
+                                                  preferredSuperview:container1
+                                                             context:nil];
+  XCTAssertTrue(dequeuedView1 == nil, @"Expected nil to be returned, since context is invalid");
+
+  UIView *dequeuedView2 = [pool dequeueStatefulViewForControllerClass:[CKOtherStatefulViewComponentController class]
+                                                   preferredSuperview:container2
+                                                              context:nil];
+  XCTAssertTrue(dequeuedView2 == view2, @"View with valid context should be returned");
+}
+
+
+- (void)testInvalidContextAfterEnqueue
+{
+  CKStatefulViewReusePool *pool = [[CKStatefulViewReusePool alloc] init];
+  CKStatefulViewComponentInvalidatableContext *context = [CKStatefulViewComponentInvalidatableContext new];
+  context.isValid = YES;
+
+  __block int calledBlockCount = 0;
+
+  UIView *container1 = [[UIView alloc] init];
+  CKTestStatefulView *view1 = [[CKTestStatefulView alloc] init];
+  [container1 addSubview:view1];
+  [pool enqueueStatefulView:view1
+         forControllerClass:[CKStatefulViewComponentWithInvalidContextController class]
+                    context:context
+         mayRelinquishBlock:^BOOL{
+           return YES;
+         }];
+  context.isValid = NO;
+
+  UIView *container2 = [[UIView alloc] init];
+  CKTestStatefulView *view2 = [[CKTestStatefulView alloc] init];
+  [container2 addSubview:view2];
+  [pool enqueueStatefulView:view2
+         forControllerClass:[CKOtherStatefulViewComponentController class]
+                    context:nil
+         mayRelinquishBlock:^BOOL{
+           calledBlockCount++;
+           return YES;
+         }];
+
+  CKRunRunLoopUntilBlockIsTrue(^BOOL{
+    return calledBlockCount == 1;
+  });
+
+  UIView *dequeuedView1 = [pool dequeueStatefulViewForControllerClass:[CKStatefulViewComponentWithInvalidContextController class]
+                                                  preferredSuperview:container1
+                                                             context:nil];
+  XCTAssertTrue(dequeuedView1 == nil, @"Expected nil to be returned, since context is invalid");
+
+  UIView *dequeuedView2 = [pool dequeueStatefulViewForControllerClass:[CKOtherStatefulViewComponentController class]
+                                                   preferredSuperview:container2
+                                                              context:nil];
+  XCTAssertTrue(dequeuedView2 == view2, @"View with valid context should be returned");
+}
+
 #pragma mark - Pending pool tests
 
 - (void)testEnqueueingViewThenDequeueingWithPendingEnabledReturnsSameViewImmediately