Fix a memory leak in SkGPipeCanvas.

src/pipe/SkGPipeWrite.cpp

Index: src/pipe/SkGPipeWrite.cpp
diff --git a/src/pipe/SkGPipeWrite.cpp b/src/pipe/SkGPipeWrite.cpp
index b61de1c912db6c3088048be23042ae405818f715..43240803256ca534e33fe31c8c1169eadeeee19c 100644
--- a/src/pipe/SkGPipeWrite.cpp
+++ b/src/pipe/SkGPipeWrite.cpp
@@ -167,32 +167,63 @@ public:
 
 ///////////////////////////////////////////////////////////////////////////////
 
+/**
+ * If SkBitmaps are to be flattened to send to the reader, this class is
+ * provided to the SkBitmapHeap to tell the SkGPipeCanvas to do so.
+ */
+class BitmapShuttle : public SkBitmapHeap::ExternalStorage {
+public:
+    BitmapShuttle(SkGPipeCanvas*);
+
+    ~BitmapShuttle();
+
+    virtual bool insert(const SkBitmap& bitmap, int32_t slot) SK_OVERRIDE;
+
+    /**
+     *  Remove the SkGPipeCanvas used for insertion. After this, calls to
+     *  insert will crash.
+     */
+    void removeCanvas();
+
+private:
+    SkGPipeCanvas*    fCanvas;
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
 class SkGPipeCanvas : public SkCanvas {
 public:
     SkGPipeCanvas(SkGPipeController*, SkWriter32*, uint32_t flags,
                   uint32_t width, uint32_t height);
     virtual ~SkGPipeCanvas();
 
-    void finish() {
-        if (!fDone) {
-            if (this->needOpBytes()) {
-                this->writeOp(kDone_DrawOp);
-                this->doNotify();
-                if (shouldFlattenBitmaps(fFlags)) {
-                    // In this case, a BitmapShuttle is reffed by the SkBitmapHeap
-                    // and refs this canvas. Unref the SkBitmapHeap to end the
-                    // circular reference. When shouldFlattenBitmaps is false,
-                    // there is no circular reference, so the SkBitmapHeap can be
-                    // safely unreffed in the destructor.
-                    fBitmapHeap->unref();
-                    // This eliminates a similar circular reference (Canvas owns
-                    // the FlattenableHeap which holds a ref to the SkBitmapHeap).
-                    fFlattenableHeap.setBitmapStorage(NULL);
-                    fBitmapHeap = NULL;
-                }
-            }
-            fDone = true;
+    /**
+     *  Called when nothing else is to be written to the stream. Any repeated
+     *  calls are ignored.
+     *
+     *  @param notifyReaders Whether to send a message to the reader(s) that
+     *      the writer is through sending commands. Should generally be true,
+     *      unless there is an error which prevents further messages from
+     *      being sent.
+     */
+    void finish(bool notifyReaders) {
+        if (fDone) {
+            return;
+        }
+        if (notifyReaders && this->needOpBytes()) {
+            this->writeOp(kDone_DrawOp);
+            this->doNotify();
         }
+        if (shouldFlattenBitmaps(fFlags)) {
+            // The following circular references exist:
+            // fFlattenableHeap -> fWriteBuffer -> fBitmapStorage -> fExternalStorage -> fCanvas
+            // fBitmapHeap -> fExternalStorage -> fCanvas
+            // fFlattenableHeap -> fBitmapStorage -> fExternalStorage -> fCanvas
+
+            // Break them all by destroying the final link to this SkGPipeCanvas.
+            fBitmapShuttle->removeCanvas();
+        }
+        fDone = true;
     }
 
     void flushRecording(bool detachCurrentBlock);
@@ -306,9 +337,11 @@ private:
     // if a new SkFlatData was added when in cross process mode
     void flattenFactoryNames();
 
-    FlattenableHeap fFlattenableHeap;
-    FlatDictionary  fFlatDictionary;
-    int fCurrFlatIndex[kCount_PaintFlats];
+    FlattenableHeap             fFlattenableHeap;
+    FlatDictionary              fFlatDictionary;
+    SkAutoTUnref<BitmapShuttle> fBitmapShuttle;
+    int                         fCurrFlatIndex[kCount_PaintFlats];
+
     int flattenToIndex(SkFlattenable* obj, PaintFlats);
 
     // Common code used by drawBitmap*. Behaves differently depending on the
@@ -390,24 +423,6 @@ int SkGPipeCanvas::flattenToIndex(SkFlattenable* obj, PaintFlats paintflat) {
 
 ///////////////////////////////////////////////////////////////////////////////
 
-/**
- * If SkBitmaps are to be flattened to send to the reader, this class is
- * provided to the SkBitmapHeap to tell the SkGPipeCanvas to do so.
- */
-class BitmapShuttle : public SkBitmapHeap::ExternalStorage {
-public:
-    BitmapShuttle(SkGPipeCanvas*);
-
-    ~BitmapShuttle();
-
-    virtual bool insert(const SkBitmap& bitmap, int32_t slot) SK_OVERRIDE;
-
-private:
-    SkGPipeCanvas*    fCanvas;
-};
-
-///////////////////////////////////////////////////////////////////////////////
-
 #define MIN_BLOCK_SIZE  (16 * 1024)
 #define BITMAPS_TO_KEEP 5
 #define FLATTENABLES_TO_KEEP 10
@@ -440,9 +455,8 @@ SkGPipeCanvas::SkGPipeCanvas(SkGPipeController* controller,
     }
 
     if (shouldFlattenBitmaps(flags)) {
-        BitmapShuttle* shuttle = SkNEW_ARGS(BitmapShuttle, (this));
-        fBitmapHeap = SkNEW_ARGS(SkBitmapHeap, (shuttle, BITMAPS_TO_KEEP));
-        shuttle->unref();
+        fBitmapShuttle.reset(SkNEW_ARGS(BitmapShuttle, (this)));
+        fBitmapHeap = SkNEW_ARGS(SkBitmapHeap, (fBitmapShuttle.get(), BITMAPS_TO_KEEP));
     } else {
         fBitmapHeap = SkNEW_ARGS(SkBitmapHeap,
                                  (BITMAPS_TO_KEEP, controller->numberOfReaders()));
@@ -456,7 +470,7 @@ SkGPipeCanvas::SkGPipeCanvas(SkGPipeController* controller,
 }
 
 SkGPipeCanvas::~SkGPipeCanvas() {
-    this->finish();
+    this->finish(true);
     SkSafeUnref(fFactorySet);
     SkSafeUnref(fBitmapHeap);
 }
@@ -474,7 +488,8 @@ bool SkGPipeCanvas::needOpBytes(size_t needed) {
         size_t blockSize = SkMax32(MIN_BLOCK_SIZE, needed);
         void* block = fController->requestBlock(blockSize, &fBlockSize);
         if (NULL == block) {
-            fDone = true;
+            // Do not notify the readers, which would call this function again.
+            this->finish(false);
             return false;
         }
         SkASSERT(SkIsAlign4(fBlockSize));
@@ -1179,7 +1194,7 @@ SkCanvas* SkGPipeWriter::startRecording(SkGPipeController* controller, uint32_t
 
 void SkGPipeWriter::endRecording() {
     if (fCanvas) {
-        fCanvas->finish();
+        fCanvas->finish(true);
         fCanvas->unref();
         fCanvas = NULL;
     }
@@ -1211,9 +1226,18 @@ BitmapShuttle::BitmapShuttle(SkGPipeCanvas* canvas) {
 }
 
 BitmapShuttle::~BitmapShuttle() {
-    fCanvas->unref();
+    this->removeCanvas();
 }
 
 bool BitmapShuttle::insert(const SkBitmap& bitmap, int32_t slot) {
+    SkASSERT(fCanvas != NULL);
     return fCanvas->shuttleBitmap(bitmap, slot);
 }
+
+void BitmapShuttle::removeCanvas() {
+    if (NULL == fCanvas) {
+        return;
+    }
+    fCanvas->unref();
+    fCanvas = NULL;
+}