{"action":"created","comment":{"url":"https://api.github.com/repos/apache/spark/pulls/comments/55898378","id":55898378,"diff_hunk":"@@ -100,48 +92,136 @@ private[spark] class MemoryStore(\n    */\n   def putBytes(blockId: BlockId, size: Long, _bytes: () => ByteBuffer): Boolean = {\n     require(!contains(blockId), s\"Block $blockId is already present in the MemoryStore\")\n-    // Work on a duplicate - since the original input might be used elsewhere.\n-    lazy val bytes = _bytes().duplicate().rewind().asInstanceOf[ByteBuffer]\n-    val putSuccess = tryToPut(blockId, () => bytes, size, deserialized = false)\n-    if (putSuccess) {\n+    if (memoryManager.acquireStorageMemory(blockId, size)) {\n+      // We acquired enough memory for the block, so go ahead and put it\n+      // Work on a duplicate - since the original input might be used elsewhere.\n+      val bytes = _bytes().duplicate().rewind().asInstanceOf[ByteBuffer]\n       assert(bytes.limit == size)\n+      val entry = new MemoryEntry(bytes, size, deserialized = false)\n+      entries.synchronized {\n+        entries.put(blockId, entry)\n+      }\n+      logInfo(\"Block %s stored as bytes in memory (estimated size %s, free %s)\".format(\n+        blockId, Utils.bytesToString(size), Utils.bytesToString(blocksMemoryUsed)))\n+      true\n+    } else {\n+      false\n     }\n-    putSuccess\n   }\n \n   /**\n    * Attempt to put the given block in memory store.\n    *\n-   * @return the estimated size of the stored data if the put() succeeded, or an iterator\n-   *         in case the put() failed (the returned iterator lets callers fall back to the disk\n-   *         store if desired).\n+   * @return in case of success, the estimated the estimated size of the stored data. In case of\n+   *         failure, return an iterator contianing the values of the block. The returned iterator\n+   *         will be backed by the combination of the partially-unrolled block and the remaining\n+   *         elements of the original input iterator. The caller must either fully consume this\n+   *         iterator or call `close()` on it in order to free the storage memory consumed by the\n+   *         partially-unrolled block.\n    */\n   private[storage] def putIterator(\n       blockId: BlockId,\n       values: Iterator[Any],\n-      level: StorageLevel): Either[Iterator[Any], Long] = {\n+      level: StorageLevel): Either[PartiallyUnrolledIterator, Long] = {\n+\n     require(!contains(blockId), s\"Block $blockId is already present in the MemoryStore\")\n-    val unrolledValues = unrollSafely(blockId, values)\n-    unrolledValues match {\n-      case Left(arrayValues) =>\n-        // Values are fully unrolled in memory, so store them as an array\n-        if (level.deserialized) {\n-          val sizeEstimate = SizeEstimator.estimate(arrayValues.asInstanceOf[AnyRef])\n-          if (tryToPut(blockId, () => arrayValues, sizeEstimate, deserialized = true)) {\n-            Right(sizeEstimate)\n-          } else {\n-            Left(arrayValues.toIterator)\n+\n+    // It's possible that the iterator is too large to materialize and store in memory. To avoid\n+    // OOM exceptions, this method will gradually unroll the iterator while periodically checking\n+    // whether there is enough free memory.\n+\n+    // Number of elements unrolled so far\n+    var elementsUnrolled = 0\n+    // Whether there is still enough memory for us to continue unrolling this block\n+    var keepUnrolling = true\n+    // Initial per-task memory to request for unrolling blocks (bytes).\n+    val initialMemoryThreshold = unrollMemoryThreshold\n+    // How often to check whether we need to request more memory\n+    val memoryCheckPeriod = 16\n+    // Memory currently reserved by this task for this particular unrolling operation\n+    var memoryThreshold = initialMemoryThreshold\n+    // Memory to request as a multiple of current vector size\n+    val memoryGrowthFactor = 1.5\n+    // Keep track of unroll memory used by this particular block / putIterator() operation\n+    var unrollMemoryUsedByThisBlock = 0L\n+    // Underlying vector for unrolling the block\n+    var vector = new SizeTrackingVector[Any]\n+\n+    // Request enough memory to begin unrolling\n+    keepUnrolling = reserveUnrollMemoryForThisTask(blockId, initialMemoryThreshold)\n+\n+    if (!keepUnrolling) {\n+      logWarning(s\"Failed to reserve initial memory threshold of \" +\n+        s\"${Utils.bytesToString(initialMemoryThreshold)} for computing block $blockId in memory.\")\n+    } else {\n+      unrollMemoryUsedByThisBlock += initialMemoryThreshold\n+    }\n+\n+    // Unroll this block safely, checking whether we have exceeded our threshold periodically\n+    while (values.hasNext && keepUnrolling) {\n+      vector += values.next()\n+      if (elementsUnrolled % memoryCheckPeriod == 0) {\n+        // If our vector's size has exceeded the threshold, request more memory\n+        val currentSize = vector.estimateSize()\n+        if (currentSize >= memoryThreshold) {\n+          val amountToRequest = (currentSize * memoryGrowthFactor - memoryThreshold).toLong\n+          keepUnrolling = reserveUnrollMemoryForThisTask(blockId, amountToRequest)\n+          if (keepUnrolling) {\n+            unrollMemoryUsedByThisBlock += amountToRequest\n           }\n+          // New threshold is currentSize * memoryGrowthFactor\n+          memoryThreshold += amountToRequest\n+        }\n+      }\n+      elementsUnrolled += 1\n+    }\n+\n+    if (keepUnrolling) {\n+      // We successfully unrolled the entirety of this block\n+      val arrayValues = vector.toArray\n+      vector = null\n+      val entry = if (level.deserialized) {\n+        val sizeEstimate = SizeEstimator.estimate(arrayValues.asInstanceOf[AnyRef])\n+        new MemoryEntry(arrayValues, sizeEstimate, deserialized = true)\n+      } else {\n+        val bytes = blockManager.dataSerialize(blockId, arrayValues.iterator)\n+        new MemoryEntry(bytes, bytes.limit, deserialized = false)\n+      }\n+      val size = entry.size\n+      val acquiredStorageMemory = {\n+        if (unrollMemoryUsedByThisBlock == size) {\n+          true\n+        } else if (unrollMemoryUsedByThisBlock < size) {\n+          memoryManager.acquireStorageMemory(blockId, size - unrollMemoryUsedByThisBlock)\n         } else {\n-          val bytes = blockManager.dataSerialize(blockId, arrayValues.iterator)\n-          if (tryToPut(blockId, () => bytes, bytes.limit, deserialized = false)) {\n-            Right(bytes.limit())\n-          } else {\n-            Left(arrayValues.toIterator)\n-          }\n+          memoryManager.releaseUnrollMemory(unrollMemoryUsedByThisBlock - size)\n+          true\n         }\n-      case Right(iteratorValues) =>\n-        Left(iteratorValues)\n+      }\n+      if (acquiredStorageMemory) {\n+        // We acquired enough memory for the block, so go ahead and put it\n+        memoryManager.synchronized {\n+          unrollMemoryMap(currentTaskAttemptId()) -= unrollMemoryUsedByThisBlock","path":"core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala","position":164,"original_position":164,"commit_id":"dbca8cf7576dc889349a70fb3c3b29c84f5a4c7c","original_commit_id":"dbca8cf7576dc889349a70fb3c3b29c84f5a4c7c","user":{"login":"andrewor14","id":2133137,"avatar_url":"https://avatars.githubusercontent.com/u/2133137?v=3","gravatar_id":"","url":"https://api.github.com/users/andrewor14","html_url":"https://github.com/andrewor14","followers_url":"https://api.github.com/users/andrewor14/followers","following_url":"https://api.github.com/users/andrewor14/following{/other_user}","gists_url":"https://api.github.com/users/andrewor14/gists{/gist_id}","starred_url":"https://api.github.com/users/andrewor14/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/andrewor14/subscriptions","organizations_url":"https://api.github.com/users/andrewor14/orgs","repos_url":"https://api.github.com/users/andrewor14/repos","events_url":"https://api.github.com/users/andrewor14/events{/privacy}","received_events_url":"https://api.github.com/users/andrewor14/received_events","type":"User","site_admin":false},"body":"can we call `releaseUnrollMemoryUsedByThisTask` instead?","created_at":"2016-03-11T22:26:45Z","updated_at":"2016-03-11T22:26:45Z","html_url":"https://github.com/apache/spark/pull/11613#discussion_r55898378","pull_request_url":"https://api.github.com/repos/apache/spark/pulls/11613","_links":{"self":{"href":"https://api.github.com/repos/apache/spark/pulls/comments/55898378"},"html":{"href":"https://github.com/apache/spark/pull/11613#discussion_r55898378"},"pull_request":{"href":"https://api.github.com/repos/apache/spark/pulls/11613"}}},"pull_request":{"url":"https://api.github.com/repos/apache/spark/pulls/11613","id":62300974,"html_url":"https://github.com/apache/spark/pull/11613","diff_url":"https://github.com/apache/spark/pull/11613.diff","patch_url":"https://github.com/apache/spark/pull/11613.patch","issue_url":"https://api.github.com/repos/apache/spark/issues/11613","number":11613,"state":"open","locked":false,"title":"[SPARK-10907][SPARK-6157] Remove pendingUnrollMemory from MemoryStore","user":{"login":"JoshRosen","id":50748,"avatar_url":"https://avatars.githubusercontent.com/u/50748?v=3","gravatar_id":"","url":"https://api.github.com/users/JoshRosen","html_url":"https://github.com/JoshRosen","followers_url":"https://api.github.com/users/JoshRosen/followers","following_url":"https://api.github.com/users/JoshRosen/following{/other_user}","gists_url":"https://api.github.com/users/JoshRosen/gists{/gist_id}","starred_url":"https://api.github.com/users/JoshRosen/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/JoshRosen/subscriptions","organizations_url":"https://api.github.com/users/JoshRosen/orgs","repos_url":"https://api.github.com/users/JoshRosen/repos","events_url":"https://api.github.com/users/JoshRosen/events{/privacy}","received_events_url":"https://api.github.com/users/JoshRosen/received_events","type":"User","site_admin":false},"body":"This patch refactors the MemoryStore to remove the concept of `pendingUnrollMemory`. It also fixes fixes SPARK-6157: \"Unrolling with MEMORY_AND_DISK should always release memory\".\r\n\r\nKey changes:\r\n\r\n- Inline `MemoryStore.tryToPut` at its three call sites in the `MemoryStore`.\r\n- Inline `Memory.unrollSafely` at its only call site (in `MemoryStore.putIterator`).\r\n- Inline `MemoryManager.acquireStorageMemory` at its call sites.\r\n- Simplify the code as a result of this inlining (some parameters have fixed values after inlining, so lots of branches can be removed).\r\n- Remove the `pendingUnrollMemory` map by returning the amount of unrollMemory allocated when returning an iterator after a failed `putIterator` call.\r\n- Change `putIterator` to return an instance of `PartiallyUnrolledIterator`, a special iterator subclass which will automatically free the unroll memory of its partially-unrolled elements when the iterator is consumed. To handle cases where the iterator is not consumed (e.g. when a MEMORY_ONLY put fails), `PartiallyUnrolledIterator` exposes a `close()` method which may be called to discard the unrolled values and free their memory.","created_at":"2016-03-09T21:42:45Z","updated_at":"2016-03-11T22:26:45Z","closed_at":null,"merged_at":null,"merge_commit_sha":"bd26f0d37b96cd7728ab96a92f48db5b82e3a890","assignee":null,"milestone":null,"commits_url":"https://api.github.com/repos/apache/spark/pulls/11613/commits","review_comments_url":"https://api.github.com/repos/apache/spark/pulls/11613/comments","review_comment_url":"https://api.github.com/repos/apache/spark/pulls/comments{/number}","comments_url":"https://api.github.com/repos/apache/spark/issues/11613/comments","statuses_url":"https://api.github.com/repos/apache/spark/statuses/dbca8cf7576dc889349a70fb3c3b29c84f5a4c7c","head":{"label":"JoshRosen:cleanup-unroll-memory","ref":"cleanup-unroll-memory","sha":"dbca8cf7576dc889349a70fb3c3b29c84f5a4c7c","user":{"login":"JoshRosen","id":50748,"avatar_url":"https://avatars.githubusercontent.com/u/50748?v=3","gravatar_id":"","url":"https://api.github.com/users/JoshRosen","html_url":"https://github.com/JoshRosen","followers_url":"https://api.github.com/users/JoshRosen/followers","following_url":"https://api.github.com/users/JoshRosen/following{/other_user}","gists_url":"https://api.github.com/users/JoshRosen/gists{/gist_id}","starred_url":"https://api.github.com/users/JoshRosen/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/JoshRosen/subscriptions","organizations_url":"https://api.github.com/users/JoshRosen/orgs","repos_url":"https://api.github.com/users/JoshRosen/repos","events_url":"https://api.github.com/users/JoshRosen/events{/privacy}","received_events_url":"https://api.github.com/users/JoshRosen/received_events","type":"User","site_admin":false},"repo":{"id":22238710,"name":"spark","full_name":"JoshRosen/spark","owner":{"login":"JoshRosen","id":50748,"avatar_url":"https://avatars.githubusercontent.com/u/50748?v=3","gravatar_id":"","url":"https://api.github.com/users/JoshRosen","html_url":"https://github.com/JoshRosen","followers_url":"https://api.github.com/users/JoshRosen/followers","following_url":"https://api.github.com/users/JoshRosen/following{/other_user}","gists_url":"https://api.github.com/users/JoshRosen/gists{/gist_id}","starred_url":"https://api.github.com/users/JoshRosen/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/JoshRosen/subscriptions","organizations_url":"https://api.github.com/users/JoshRosen/orgs","repos_url":"https://api.github.com/users/JoshRosen/repos","events_url":"https://api.github.com/users/JoshRosen/events{/privacy}","received_events_url":"https://api.github.com/users/JoshRosen/received_events","type":"User","site_admin":false},"private":false,"html_url":"https://github.com/JoshRosen/spark","description":"Mirror of Apache Spark","fork":true,"url":"https://api.github.com/repos/JoshRosen/spark","forks_url":"https://api.github.com/repos/JoshRosen/spark/forks","keys_url":"https://api.github.com/repos/JoshRosen/spark/keys{/key_id}","collaborators_url":"https://api.github.com/repos/JoshRosen/spark/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/JoshRosen/spark/teams","hooks_url":"https://api.github.com/repos/JoshRosen/spark/hooks","issue_events_url":"https://api.github.com/repos/JoshRosen/spark/issues/events{/number}","events_url":"https://api.github.com/repos/JoshRosen/spark/events","assignees_url":"https://api.github.com/repos/JoshRosen/spark/assignees{/user}","branches_url":"https://api.github.com/repos/JoshRosen/spark/branches{/branch}","tags_url":"https://api.github.com/repos/JoshRosen/spark/tags","blobs_url":"https://api.github.com/repos/JoshRosen/spark/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/JoshRosen/spark/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/JoshRosen/spark/git/refs{/sha}","trees_url":"https://api.github.com/repos/JoshRosen/spark/git/trees{/sha}","statuses_url":"https://api.github.com/repos/JoshRosen/spark/statuses/{sha}","languages_url":"https://api.github.com/repos/JoshRosen/spark/languages","stargazers_url":"https://api.github.com/repos/JoshRosen/spark/stargazers","contributors_url":"https://api.github.com/repos/JoshRosen/spark/contributors","subscribers_url":"https://api.github.com/repos/JoshRosen/spark/subscribers","subscription_url":"https://api.github.com/repos/JoshRosen/spark/subscription","commits_url":"https://api.github.com/repos/JoshRosen/spark/commits{/sha}","git_commits_url":"https://api.github.com/repos/JoshRosen/spark/git/commits{/sha}","comments_url":"https://api.github.com/repos/JoshRosen/spark/comments{/number}","issue_comment_url":"https://api.github.com/repos/JoshRosen/spark/issues/comments{/number}","contents_url":"https://api.github.com/repos/JoshRosen/spark/contents/{+path}","compare_url":"https://api.github.com/repos/JoshRosen/spark/compare/{base}...{head}","merges_url":"https://api.github.com/repos/JoshRosen/spark/merges","archive_url":"https://api.github.com/repos/JoshRosen/spark/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/JoshRosen/spark/downloads","issues_url":"https://api.github.com/repos/JoshRosen/spark/issues{/number}","pulls_url":"https://api.github.com/repos/JoshRosen/spark/pulls{/number}","milestones_url":"https://api.github.com/repos/JoshRosen/spark/milestones{/number}","notifications_url":"https://api.github.com/repos/JoshRosen/spark/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/JoshRosen/spark/labels{/name}","releases_url":"https://api.github.com/repos/JoshRosen/spark/releases{/id}","deployments_url":"https://api.github.com/repos/JoshRosen/spark/deployments","created_at":"2014-07-25T00:47:52Z","updated_at":"2015-05-01T22:45:32Z","pushed_at":"2016-03-11T21:48:29Z","git_url":"git://github.com/JoshRosen/spark.git","ssh_url":"git@github.com:JoshRosen/spark.git","clone_url":"https://github.com/JoshRosen/spark.git","svn_url":"https://github.com/JoshRosen/spark","homepage":null,"size":153121,"stargazers_count":1,"watchers_count":1,"language":"Scala","has_issues":false,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":3,"mirror_url":null,"open_issues_count":0,"forks":3,"open_issues":0,"watchers":1,"default_branch":"master"}},"base":{"label":"apache:master","ref":"master","sha":"073bf9d4d91e0242a813f3d227e52e76c26a2200","user":{"login":"apache","id":47359,"avatar_url":"https://avatars.githubusercontent.com/u/47359?v=3","gravatar_id":"","url":"https://api.github.com/users/apache","html_url":"https://github.com/apache","followers_url":"https://api.github.com/users/apache/followers","following_url":"https://api.github.com/users/apache/following{/other_user}","gists_url":"https://api.github.com/users/apache/gists{/gist_id}","starred_url":"https://api.github.com/users/apache/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/apache/subscriptions","organizations_url":"https://api.github.com/users/apache/orgs","repos_url":"https://api.github.com/users/apache/repos","events_url":"https://api.github.com/users/apache/events{/privacy}","received_events_url":"https://api.github.com/users/apache/received_events","type":"Organization","site_admin":false},"repo":{"id":17165658,"name":"spark","full_name":"apache/spark","owner":{"login":"apache","id":47359,"avatar_url":"https://avatars.githubusercontent.com/u/47359?v=3","gravatar_id":"","url":"https://api.github.com/users/apache","html_url":"https://github.com/apache","followers_url":"https://api.github.com/users/apache/followers","following_url":"https://api.github.com/users/apache/following{/other_user}","gists_url":"https://api.github.com/users/apache/gists{/gist_id}","starred_url":"https://api.github.com/users/apache/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/apache/subscriptions","organizations_url":"https://api.github.com/users/apache/orgs","repos_url":"https://api.github.com/users/apache/repos","events_url":"https://api.github.com/users/apache/events{/privacy}","received_events_url":"https://api.github.com/users/apache/received_events","type":"Organization","site_admin":false},"private":false,"html_url":"https://github.com/apache/spark","description":"Mirror of Apache Spark","fork":false,"url":"https://api.github.com/repos/apache/spark","forks_url":"https://api.github.com/repos/apache/spark/forks","keys_url":"https://api.github.com/repos/apache/spark/keys{/key_id}","collaborators_url":"https://api.github.com/repos/apache/spark/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/apache/spark/teams","hooks_url":"https://api.github.com/repos/apache/spark/hooks","issue_events_url":"https://api.github.com/repos/apache/spark/issues/events{/number}","events_url":"https://api.github.com/repos/apache/spark/events","assignees_url":"https://api.github.com/repos/apache/spark/assignees{/user}","branches_url":"https://api.github.com/repos/apache/spark/branches{/branch}","tags_url":"https://api.github.com/repos/apache/spark/tags","blobs_url":"https://api.github.com/repos/apache/spark/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/apache/spark/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/apache/spark/git/refs{/sha}","trees_url":"https://api.github.com/repos/apache/spark/git/trees{/sha}","statuses_url":"https://api.github.com/repos/apache/spark/statuses/{sha}","languages_url":"https://api.github.com/repos/apache/spark/languages","stargazers_url":"https://api.github.com/repos/apache/spark/stargazers","contributors_url":"https://api.github.com/repos/apache/spark/contributors","subscribers_url":"https://api.github.com/repos/apache/spark/subscribers","subscription_url":"https://api.github.com/repos/apache/spark/subscription","commits_url":"https://api.github.com/repos/apache/spark/commits{/sha}","git_commits_url":"https://api.github.com/repos/apache/spark/git/commits{/sha}","comments_url":"https://api.github.com/repos/apache/spark/comments{/number}","issue_comment_url":"https://api.github.com/repos/apache/spark/issues/comments{/number}","contents_url":"https://api.github.com/repos/apache/spark/contents/{+path}","compare_url":"https://api.github.com/repos/apache/spark/compare/{base}...{head}","merges_url":"https://api.github.com/repos/apache/spark/merges","archive_url":"https://api.github.com/repos/apache/spark/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/apache/spark/downloads","issues_url":"https://api.github.com/repos/apache/spark/issues{/number}","pulls_url":"https://api.github.com/repos/apache/spark/pulls{/number}","milestones_url":"https://api.github.com/repos/apache/spark/milestones{/number}","notifications_url":"https://api.github.com/repos/apache/spark/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/apache/spark/labels{/name}","releases_url":"https://api.github.com/repos/apache/spark/releases{/id}","deployments_url":"https://api.github.com/repos/apache/spark/deployments","created_at":"2014-02-25T08:00:08Z","updated_at":"2016-03-11T22:23:12Z","pushed_at":"2016-03-11T22:24:10Z","git_url":"git://github.com/apache/spark.git","ssh_url":"git@github.com:apache/spark.git","clone_url":"https://github.com/apache/spark.git","svn_url":"https://github.com/apache/spark","homepage":null,"size":440207,"stargazers_count":7769,"watchers_count":7769,"language":"Scala","has_issues":false,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":7086,"mirror_url":"git://git.apache.org/spark.git","open_issues_count":391,"forks":7086,"open_issues":391,"watchers":7769,"default_branch":"master"}},"_links":{"self":{"href":"https://api.github.com/repos/apache/spark/pulls/11613"},"html":{"href":"https://github.com/apache/spark/pull/11613"},"issue":{"href":"https://api.github.com/repos/apache/spark/issues/11613"},"comments":{"href":"https://api.github.com/repos/apache/spark/issues/11613/comments"},"review_comments":{"href":"https://api.github.com/repos/apache/spark/pulls/11613/comments"},"review_comment":{"href":"https://api.github.com/repos/apache/spark/pulls/comments{/number}"},"commits":{"href":"https://api.github.com/repos/apache/spark/pulls/11613/commits"},"statuses":{"href":"https://api.github.com/repos/apache/spark/statuses/dbca8cf7576dc889349a70fb3c3b29c84f5a4c7c"}}}}