SkPDF: wait for jobs to complete on abort()

Before, we could only wait for all reserved objects to serialize.

The new unit test `SkPDF_abort_jobs` crashed before this change.

Change-Id: Ia2fcb33251c6c32c125f631ed3eeb619f616bef3
Reviewed-on: https://skia-review.googlesource.com/c/179856
Auto-Submit: Hal Canary <halcanary@google.com>
Commit-Queue: Herb Derby <herb@google.com>
Reviewed-by: Herb Derby <herb@google.com>
diff --git a/src/pdf/SkPDFBitmap.cpp b/src/pdf/SkPDFBitmap.cpp
index 7fc9c4e..4f7ae95 100644
--- a/src/pdf/SkPDFBitmap.cpp
+++ b/src/pdf/SkPDFBitmap.cpp
@@ -310,9 +310,11 @@
     SkPDFIndirectReference ref = doc->reserveRef();
     if (SkExecutor* executor = doc->executor()) {
         SkRef(img);
+        doc->incrementJobCount();
         executor->add([img, encodingQuality, doc, ref]() {
             serialize_image(img, encodingQuality, doc, ref);
             SkSafeUnref(img);
+            doc->signalJobComplete();
         });
         return ref;
     }
diff --git a/src/pdf/SkPDFDocument.cpp b/src/pdf/SkPDFDocument.cpp
index a1fd985..9cceda5 100644
--- a/src/pdf/SkPDFDocument.cpp
+++ b/src/pdf/SkPDFDocument.cpp
@@ -218,7 +218,6 @@
 void SkPDFDocument::endObject() {
     end_indirect_object(this->getStream());
     fMutex.release();
-    fSemaphore.signal();
 };
 
 static SkSize operator*(SkISize u, SkScalar s) { return SkSize{u.width() * s, u.height() * s}; }
@@ -292,6 +291,7 @@
 }
 
 void SkPDFDocument::onAbort() {
+    this->waitForJobs();
     this->reset();
 }
 
@@ -468,6 +468,7 @@
 void SkPDFDocument::onClose(SkWStream* stream) {
     SkASSERT(fCanvas.imageInfo().dimensions().isZero());
     if (fPages.empty()) {
+        this->waitForJobs();
         this->reset();
         return;
     }
@@ -502,13 +503,7 @@
         f->emitSubset(this);
     }
 
-    int waits = 0;
-    // fNextObjectNumber can increase while we wait.
-    while (waits + 1 < fNextObjectNumber.load()) {
-        fSemaphore.wait();
-        ++waits;
-    }
-    SkASSERT(fNextObjectNumber.load() == fOffsetMap.objectCount());
+    this->waitForJobs();
     {
         SkAutoMutexAcquire autoMutexAcquire(fMutex);
         serialize_footer(fOffsetMap, this->getStream(), fInfoDict, docCatalogRef, fUUID);
@@ -516,6 +511,17 @@
     this->reset();
 }
 
+void SkPDFDocument::incrementJobCount() { fJobCount++; }
+
+void SkPDFDocument::signalJobComplete() { fSemaphore.signal(); }
+
+void SkPDFDocument::waitForJobs() {
+     // fJobCount can increase while we wait.
+     while (fJobCount > 0) {
+         fSemaphore.wait();
+         --fJobCount;
+     }
+}
 
 ///////////////////////////////////////////////////////////////////////////////
 
diff --git a/src/pdf/SkPDFDocumentPriv.h b/src/pdf/SkPDFDocumentPriv.h
index 88f3c2a..58799e6 100644
--- a/src/pdf/SkPDFDocumentPriv.h
+++ b/src/pdf/SkPDFDocumentPriv.h
@@ -71,6 +71,8 @@
     void endObject();
 
     SkExecutor* executor() const { return fExecutor; }
+    void incrementJobCount();
+    void signalJobComplete();
     size_t currentPageIndex() { return fPages.size(); }
     size_t pageCount() { return fPageRefs.size(); }
 
@@ -83,6 +85,7 @@
     SkPDFDict fDests;
     sk_sp<SkPDFDevice> fPageDevice;
     std::atomic<int> fNextObjectNumber = {1};
+    std::atomic<int> fJobCount = {0};
     SkUUID fUUID;
     SkPDFIndirectReference fInfoDict;
     SkPDFIndirectReference fXMP;
@@ -98,6 +101,7 @@
     SkSemaphore fSemaphore;
 
     void reset();
+    void waitForJobs();
 };
 
 #endif  // SkPDFDocumentPriv_DEFINED
diff --git a/src/pdf/SkPDFTypes.cpp b/src/pdf/SkPDFTypes.cpp
index 0b98fca..76eb58e 100644
--- a/src/pdf/SkPDFTypes.cpp
+++ b/src/pdf/SkPDFTypes.cpp
@@ -484,10 +484,12 @@
         SkStreamAsset* contentPtr = content.release();
         // Pass ownership of both pointers into a std::function, which should
         // only be executed once.
+        doc->incrementJobCount();
         executor->add([dictPtr, contentPtr, deflate, doc, ref]() {
             serialize_stream(dictPtr, contentPtr, deflate, doc, ref);
             delete dictPtr;
             delete contentPtr;
+            doc->signalJobComplete();
         });
         return ref;
     }
diff --git a/tests/PDFDocumentTest.cpp b/tests/PDFDocumentTest.cpp
index b4941dc..bf1fedf 100644
--- a/tests/PDFDocumentTest.cpp
+++ b/tests/PDFDocumentTest.cpp
@@ -8,6 +8,7 @@
 
 #include "Resources.h"
 #include "SkCanvas.h"
+#include "SkExecutor.h"
 #include "SkOSFile.h"
 #include "SkOSPath.h"
 #include "SkPDFDocument.h"
@@ -239,3 +240,19 @@
                 SkColorSetARGB(0xFF, 0x00, (uint8_t)(255.0f * i / (n - 1)), 0x00));
     }
 }
+
+// Test to make sure that jobs launched by PDF backend don't cause a segfault
+// after calling abort().
+DEF_TEST(SkPDF_abort_jobs, rep) {
+    SkBitmap b;
+    b.allocN32Pixels(612, 792);
+    b.eraseColor(0x4F9643A0);
+    SkPDF::Metadata metadata;
+    std::unique_ptr<SkExecutor> executor = SkExecutor::MakeFIFOThreadPool();
+    metadata.fExecutor = executor.get();
+    SkNullWStream dst;
+    sk_sp<SkDocument> doc = SkPDF::MakeDocument(&dst, metadata);
+    doc->beginPage(612, 792)->drawBitmap(b, 0, 0);
+    doc->abort();
+}
+