Update reshares in batches as not to run out of memory

2 files changed, 97 insertions, 14 deletions

diff --git a/apps/files_sharing/lib/migration.php b/apps/files_sharing/lib/migration.php
index 91e0c4c0996..4ec411d81da 100644
--- a/apps/files_sharing/lib/migration.php
+++ b/apps/files_sharing/lib/migration.php
@@ -23,6 +23,7 @@ namespace OCA\Files_Sharing;
 
 use Doctrine\DBAL\Connection;
 use OCP\IDBConnection;
+use OC\Cache\CappedMemoryCache;
 
 /**
  * Class Migration
@@ -43,6 +44,9 @@ class Migration {
 
 	public function __construct(IDBConnection $connection) {
 		$this->connection = $connection;
+
+		// We cache up to 10k share items (~20MB)
+		$this->shareCache = new CappedMemoryCache(10000);
 	}
 
 	/**
@@ -50,17 +54,33 @@ class Migration {
 	 * upgrade from oC 8.2 to 9.0 with the new sharing
 	 */
 	public function removeReShares() {
-		$reShares = $this->getAllReShares();
-		$this->shareCache = $reShares;
-		$owners = [];
-		foreach ($reShares as $share) {
-			$owners[$share['id']] = [
-				'owner' => $this->findOwner($share),
-				'initiator' => $share['uid_owner']
-			];
-		}
 
-		$this->updateOwners($owners);
+		while(true) {
+			$reShares = $this->getReShares(1000);
+
+			if (empty($reShares)) {
+				break;
+			}
+
+			// Update the cache
+			foreach($reShares as $reShare) {
+				$this->shareCache[$reShare['id']] = $reShare;
+			}
+
+			$owners = [];
+			foreach ($reShares as $share) {
+				$owners[$share['id']] = [
+					'owner' => $this->findOwner($share),
+					'initiator' => $share['uid_owner']
+				];
+			}
+			$this->updateOwners($owners);
+
+			//Clear the cache of the shares we just updated so we have more room
+			foreach($owners as $id => $owner) {
+				unset($this->shareCache[$id]);
+			}
+		}
 	}
 
 	/**
@@ -84,11 +104,12 @@ class Migration {
 	}
 
 	/**
-	 * get all re-shares from the database
+	 * Get $n re-shares from the database
 	 *
+	 * @param int $n The max number of shares to fetch
 	 * @return array
 	 */
-	private function getAllReShares() {
+	private function getReShares($n = 1000) {
 		$query = $this->connection->getQueryBuilder();
 		$query->select(['id', 'parent', 'uid_owner'])
 			->from($this->table)
@@ -110,7 +131,9 @@ class Migration {
 					Connection::PARAM_STR_ARRAY
 				)
 			))
-			->andWhere($query->expr()->isNotNull('parent'));
+			->andWhere($query->expr()->isNotNull('parent'))
+			->orderBy('id', 'asc')
+			->setMaxResults($n);
 		$result = $query->execute();
 		$shares = $result->fetchAll();
 		$result->closeCursor();
@@ -159,7 +182,8 @@ class Migration {
 					->set('parent', $query->createNamedParameter(null))
 					->set('uid_owner', $query->createNamedParameter($owner['owner']))
 					->set('uid_initiator', $query->createNamedParameter($owner['initiator']))
-					->where($query->expr()->eq('id', $query->createNamedParameter($id)))->execute();
+					->where($query->expr()->eq('id', $query->createNamedParameter($id)))
+					->execute();
 			}
 
 			$this->connection->commit();
diff --git a/apps/files_sharing/tests/migrationtest.php b/apps/files_sharing/tests/migrationtest.php
index 6b0e5e3ca29..fb9d1242ac0 100644
--- a/apps/files_sharing/tests/migrationtest.php
+++ b/apps/files_sharing/tests/migrationtest.php
@@ -246,4 +246,63 @@ class MigrationTest extends TestCase {
 		}
 	}
 
+	public function test100kDeepReshares() {
+		$parent = null;
+		for ($i = 0; $i < 10; $i++) {
+			$query = $this->connection->getQueryBuilder();
+			$query->insert($this->table)
+				->values(
+					[
+						'share_type' => $query->createParameter('share_type'),
+						'share_with' => $query->createParameter('share_with'),
+						'uid_owner' => $query->createParameter('uid_owner'),
+						'uid_initiator' => $query->createParameter('uid_initiator'),
+						'parent' => $query->createParameter('parent'),
+						'item_type' => $query->createParameter('item_type'),
+						'item_source' => $query->createParameter('item_source'),
+						'item_target' => $query->createParameter('item_target'),
+						'file_source' => $query->createParameter('file_source'),
+						'file_target' => $query->createParameter('file_target'),
+						'permissions' => $query->createParameter('permissions'),
+						'stime' => $query->createParameter('stime'),
+					]
+				)
+				->setParameter('share_type', \OCP\Share::SHARE_TYPE_USER)
+				->setParameter('share_with', 'user'.($i+1))
+				->setParameter('uid_owner', 'user'.($i))
+				->setParameter('uid_initiator', '')
+				->setParameter('parent', $parent)
+				->setParameter('item_type', 'file')
+				->setParameter('item_source', '2')
+				->setParameter('item_target', '/2')
+				->setParameter('file_source', 2)
+				->setParameter('file_target', '/foobar')
+				->setParameter('permissions', 31)
+				->setParameter('stime', time());
+
+			$this->assertSame(1, $query->execute());
+			$parent = $query->getLastInsertId();
+		}
+
+		$this->migration->removeReShares();
+
+		$qb = $this->connection->getQueryBuilder();
+
+		$stmt = $qb->select('id', 'share_with', 'uid_owner', 'uid_initiator', 'parent')
+			->from('share')
+			->orderBy('id', 'asc')
+			->execute();
+
+		$i = 0;
+		while($share = $stmt->fetch()) {
+			$this->assertEquals('user'.($i+1), $share['share_with']);
+			if ($i !== 0) {
+				$this->assertEquals('user' . ($i), $share['uid_initiator']);
+				$this->assertEquals('user0', $share['uid_owner']);
+			}
+			$this->assertEquals(null, $share['parent']);
+			$i++;
+		}
+		$stmt->closeCursor();
+	}
 }