Revert "Revert "Add method to sk_gpu_test::TestContext to automatically restore the previous context.""

This reverts commit 1e09e461d2ffcf8b07242cfe93dd7d12c4d75866.

Change-Id: I95d5544a7baaa078536790493ce4119816a77e94
Reviewed-on: https://skia-review.googlesource.com/72903
Reviewed-by: Brian Osman <brianosman@google.com>
Commit-Queue: Brian Salomon <bsalomon@google.com>
diff --git a/gn/core.gni b/gn/core.gni
index 358c387..186e5af 100644
--- a/gn/core.gni
+++ b/gn/core.gni
@@ -272,6 +272,7 @@
   "$_src/core/SkScan_Antihair.cpp",
   "$_src/core/SkScan_Hairline.cpp",
   "$_src/core/SkScan_Path.cpp",
+  "$_src/core/SkScopeExit.h",
   "$_src/core/SkSemaphore.cpp",
   "$_src/core/SkSharedMutex.cpp",
   "$_src/core/SkSharedMutex.h",
diff --git a/gn/pdf.gni b/gn/pdf.gni
index 355c27a..665e35c 100644
--- a/gn/pdf.gni
+++ b/gn/pdf.gni
@@ -46,5 +46,4 @@
   "$_src/pdf/SkPDFTypes.h",
   "$_src/pdf/SkPDFUtils.cpp",
   "$_src/pdf/SkPDFUtils.h",
-  "$_src/pdf/SkScopeExit.h",
 ]
diff --git a/src/core/SkScopeExit.h b/src/core/SkScopeExit.h
new file mode 100644
index 0000000..95804e6
--- /dev/null
+++ b/src/core/SkScopeExit.h
@@ -0,0 +1,53 @@
+/*
+ * Copyright 2016 Google Inc.
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#ifndef SkScopeExit_DEFINED
+#define SkScopeExit_DEFINED
+
+#include "SkTypes.h"
+
+/** SkScopeExit calls a std:::function<void()> in its destructor. */
+class SkScopeExit {
+public:
+    SkScopeExit(std::function<void()> f) : fFn(std::move(f)) {}
+    SkScopeExit(SkScopeExit&& that) : fFn(std::move(that.fFn)) {}
+
+    ~SkScopeExit() {
+        if (fFn) {
+            fFn();
+        }
+    }
+
+    SkScopeExit& operator=(SkScopeExit&& that) {
+        fFn = std::move(that.fFn);
+        return *this;
+    }
+
+private:
+    std::function<void()> fFn;
+
+    SkScopeExit(           const SkScopeExit& ) = delete;
+    SkScopeExit& operator=(const SkScopeExit& ) = delete;
+};
+
+/**
+ * SK_AT_SCOPE_EXIT(stmt) evaluates stmt when the current scope ends.
+ *
+ * E.g.
+ *    {
+ *        int x = 5;
+ *        {
+ *           SK_AT_SCOPE_EXIT(x--);
+ *           SkASSERT(x == 5);
+ *        }
+ *        SkASSERT(x == 4);
+ *    }
+ */
+#define SK_AT_SCOPE_EXIT(stmt)                              \
+    SkScopeExit SK_MACRO_APPEND_LINE(at_scope_exit_)([&]() { stmt; })
+
+#endif  // SkScopeExit_DEFINED
diff --git a/src/pdf/SkScopeExit.h b/src/pdf/SkScopeExit.h
deleted file mode 100644
index 5b7bcdc..0000000
--- a/src/pdf/SkScopeExit.h
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright 2016 Google Inc.
- *
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- */
-
-#ifndef SkScopeExit_DEFINED
-#define SkScopeExit_DEFINED
-
-#include "SkTypes.h"
-
-/**
- * SK_AT_SCOPE_EXIT(stmt) evaluates stmt when the current scope ends.
- *
- * E.g.
- *    {
- *        int x = 5;
- *        {
- *           SK_AT_SCOPE_EXIT(x--);
- *           SkASSERT(x == 5);
- *        }
- *        SkASSERT(x == 4);
- *    }
- */
-template <typename Fn>
-class SkScopeExit {
-public:
-    SkScopeExit(Fn f) : fFn(std::move(f)) {}
-    ~SkScopeExit() { fFn(); }
-
-private:
-    Fn fFn;
-
-    SkScopeExit(           const SkScopeExit& ) = delete;
-    SkScopeExit& operator=(const SkScopeExit& ) = delete;
-    SkScopeExit(                 SkScopeExit&&) = delete;
-    SkScopeExit& operator=(      SkScopeExit&&) = delete;
-};
-
-template <typename Fn>
-inline SkScopeExit<Fn> SkMakeScopeExit(Fn&& fn) {
-    return {std::move(fn)};
-}
-
-#define SK_AT_SCOPE_EXIT(stmt)                              \
-    SK_UNUSED auto&& SK_MACRO_APPEND_LINE(at_scope_exit_) = \
-        SkMakeScopeExit([&]() { stmt; });
-
-#endif  // SkScopeExit_DEFINED
diff --git a/tests/ImageTest.cpp b/tests/ImageTest.cpp
index 30ae635..d18e6a9 100644
--- a/tests/ImageTest.cpp
+++ b/tests/ImageTest.cpp
@@ -436,10 +436,10 @@
         // Create a texture image.
         [context] { return create_gpu_image(context); },
         // Create a texture image in a another GrContext.
-        [testContext, otherContextInfo] {
-            otherContextInfo.testContext()->makeCurrent();
+        [otherContextInfo] {
+            auto restore = otherContextInfo.testContext()->makeCurrentAndAutoRestore();
             sk_sp<SkImage> otherContextImage = create_gpu_image(otherContextInfo.grContext());
-            testContext->makeCurrent();
+            otherContextInfo.grContext()->flush();
             return otherContextImage;
         }
     };
@@ -487,7 +487,6 @@
             }
         }
 
-        testContext->makeCurrent();
         context->flush();
     }
 }
@@ -1197,10 +1196,10 @@
         { create_picture_image, true, false },
         { [context] { return create_gpu_image(context); }, true, true },
         // Create a texture image in a another GrContext.
-        { [testContext, otherContextInfo] {
-            otherContextInfo.testContext()->makeCurrent();
+        { [otherContextInfo] {
+            auto restore = otherContextInfo.testContext()->makeCurrentAndAutoRestore();
             sk_sp<SkImage> otherContextImage = create_gpu_image(otherContextInfo.grContext());
-            testContext->makeCurrent();
+            otherContextInfo.grContext()->flush();
             return otherContextImage;
           }, false, false },
         // Create an image that is too large to be texture backed.
@@ -1232,7 +1231,6 @@
             kExpectedState[testCase.fCanTakeDirectly]);
         }
 
-        testContext->makeCurrent();
         context->flush();
     }
 }
diff --git a/tools/gpu/GrContextFactory.cpp b/tools/gpu/GrContextFactory.cpp
index 21766db..00c4b65 100644
--- a/tools/gpu/GrContextFactory.cpp
+++ b/tools/gpu/GrContextFactory.cpp
@@ -52,8 +52,9 @@
 
 void GrContextFactory::destroyContexts() {
     for (Context& context : fContexts) {
+        SkScopeExit restore(nullptr);
         if (context.fTestContext) {
-            context.fTestContext->makeCurrent();
+            restore = context.fTestContext->makeCurrentAndAutoRestore();
         }
         if (!context.fGrContext->unique()) {
             context.fGrContext->releaseResourcesAndAbandonContext();
@@ -69,7 +70,7 @@
     for (Context& context : fContexts) {
         if (!context.fAbandoned) {
             if (context.fTestContext) {
-                context.fTestContext->makeCurrent();
+                auto restore = context.fTestContext->makeCurrentAndAutoRestore();
                 context.fTestContext->testAbandon();
                 delete(context.fTestContext);
                 context.fTestContext = nullptr;
@@ -82,9 +83,10 @@
 
 void GrContextFactory::releaseResourcesAndAbandonContexts() {
     for (Context& context : fContexts) {
+        SkScopeExit restore(nullptr);
         if (!context.fAbandoned) {
             if (context.fTestContext) {
-                context.fTestContext->makeCurrent();
+                restore = context.fTestContext->makeCurrentAndAutoRestore();
             }
             context.fGrContext->releaseResourcesAndAbandonContext();
             context.fAbandoned = true;
@@ -237,7 +239,7 @@
         default:
             return ContextInfo();
     }
-    testCtx->makeCurrent();
+
     SkASSERT(testCtx && testCtx->backend() == backend);
     GrContextOptions grOptions = fGlobalOptions;
     if (ContextOverrides::kDisableNVPR & overrides) {
@@ -252,7 +254,11 @@
     if (ContextOverrides::kAvoidStencilBuffers & overrides) {
         grOptions.fAvoidStencilBuffers = true;
     }
-    sk_sp<GrContext> grCtx = testCtx->makeGrContext(grOptions);
+    sk_sp<GrContext> grCtx;
+    {
+        auto restore = testCtx->makeCurrentAndAutoRestore();
+        grCtx = testCtx->makeGrContext(grOptions);
+    }
     if (!grCtx.get()) {
         return ContextInfo();
     }
@@ -282,6 +288,7 @@
     context.fShareContext = shareContext;
     context.fShareIndex = shareIndex;
     context.fOptions = grOptions;
+    context.fTestContext->makeCurrent();
     return ContextInfo(context.fType, context.fTestContext, context.fGrContext, context.fOptions);
 }
 
diff --git a/tools/gpu/TestContext.cpp b/tools/gpu/TestContext.cpp
index c80c4ea..f760cc7 100644
--- a/tools/gpu/TestContext.cpp
+++ b/tools/gpu/TestContext.cpp
@@ -37,8 +37,15 @@
 
 void TestContext::makeCurrent() const { this->onPlatformMakeCurrent(); }
 
+SkScopeExit TestContext::makeCurrentAndAutoRestore() const {
+    auto asr = SkScopeExit(this->onPlatformGetAutoContextRestore());
+    this->makeCurrent();
+    return asr;
+}
+
 void TestContext::swapBuffers() { this->onPlatformSwapBuffers(); }
 
+
 void TestContext::waitOnSyncOrSwap() {
     if (!fFenceSync) {
         // Fallback on the platform SwapBuffers method for synchronization. This may have no effect.
diff --git a/tools/gpu/TestContext.h b/tools/gpu/TestContext.h
index 84794f3..5b512db 100644
--- a/tools/gpu/TestContext.h
+++ b/tools/gpu/TestContext.h
@@ -9,10 +9,11 @@
 #ifndef TestContext_DEFINED
 #define TestContext_DEFINED
 
+#include "../private/SkTemplates.h"
 #include "FenceSync.h"
 #include "GrTypes.h"
 #include "SkRefCnt.h"
-#include "../private/SkTemplates.h"
+#include "SkScopeExit.h"
 
 class GrContext;
 struct GrContextOptions;
@@ -45,6 +46,18 @@
 
     void makeCurrent() const;
 
+    /**
+     * Like makeCurrent() but this returns an object that will restore the previous current
+     * context in its destructor. Useful to undo the effect making this current before returning to
+     * a caller that doesn't expect the current context to be changed underneath it.
+     *
+     * The returned object restores the current context of the same type (e.g. egl, glx, ...) in its
+     * destructor. It is undefined behavior if that context is destroyed before the destructor
+     * executes. If the concept of a current context doesn't make sense for this context type then
+     * the returned object's destructor is a no-op.
+     */
+    SkScopeExit SK_WARN_UNUSED_RESULT makeCurrentAndAutoRestore() const;
+
     virtual GrBackend backend() = 0;
     virtual GrBackendContext backendContext() = 0;
 
@@ -94,6 +107,14 @@
     virtual void teardown();
 
     virtual void onPlatformMakeCurrent() const = 0;
+    /**
+     * Subclasses should implement such that the returned function will cause the current context
+     * of this type to be made current again when it is called. It should additionally be the
+     * case that if "this" is already current when this is called, then "this" is destroyed (thereby
+     * setting the null context as current), and then the std::function is called the null context
+     * should remain current.
+     */
+    virtual std::function<void()> onPlatformGetAutoContextRestore() const = 0;
     virtual void onPlatformSwapBuffers() const = 0;
 
 private:
diff --git a/tools/gpu/gl/angle/GLTestContext_angle.cpp b/tools/gpu/gl/angle/GLTestContext_angle.cpp
index 52cc512..3b55c40 100644
--- a/tools/gpu/gl/angle/GLTestContext_angle.cpp
+++ b/tools/gpu/gl/angle/GLTestContext_angle.cpp
@@ -34,6 +34,16 @@
     void* fEGLLib;
 };
 
+std::function<void()> context_restorer() {
+    auto display = eglGetCurrentDisplay();
+    auto dsurface = eglGetCurrentSurface(EGL_DRAW);
+    auto rsurface = eglGetCurrentSurface(EGL_READ);
+    auto context = eglGetCurrentContext();
+    return [display, dsurface, rsurface, context] {
+        eglMakeCurrent(display, dsurface, rsurface, context);
+    };
+}
+
 static GrGLFuncPtr angle_get_gl_proc(void* ctx, const char name[]) {
     const Libs* libs = reinterpret_cast<const Libs*>(ctx);
     GrGLFuncPtr proc = (GrGLFuncPtr) GetProcedureAddress(libs->fGLLib, name);
@@ -87,6 +97,7 @@
     void destroyGLContext();
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char* name) const override;
 
@@ -214,6 +225,7 @@
 
     fSurface = eglCreatePbufferSurface(fDisplay, surfaceConfig, surfaceAttribs);
 
+    SkScopeExit restorer(context_restorer());
     if (!eglMakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
         SkDebugf("Could not set the context.");
         this->destroyGLContext();
@@ -320,7 +332,10 @@
 
 void ANGLEGLContext::destroyGLContext() {
     if (EGL_NO_DISPLAY != fDisplay) {
-        eglMakeCurrent(fDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+        if (eglGetCurrentContext() == fContext) {
+            // This will ensure that the context is immediately deleted.
+            eglMakeCurrent(fDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+        }
 
         if (EGL_NO_CONTEXT != fContext) {
             eglDestroyContext(fDisplay, fContext);
@@ -355,6 +370,13 @@
     }
 }
 
+std::function<void()> ANGLEGLContext::onPlatformGetAutoContextRestore() const {
+    if (eglGetCurrentContext() == fContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void ANGLEGLContext::onPlatformSwapBuffers() const {
     if (!eglSwapBuffers(fDisplay, fSurface)) {
         SkDebugf("Could not complete eglSwapBuffers.\n");
diff --git a/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.cpp b/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.cpp
index 54845fc..be2b6ad 100644
--- a/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.cpp
+++ b/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.cpp
@@ -10,11 +10,14 @@
 
 #include "SkMutex.h"
 #include "SkOnce.h"
+#include "SkTLS.h"
 #include "gl/GrGLInterface.h"
 #include "gl/GrGLAssembleInterface.h"
 #include "gl/command_buffer/GLTestContext_command_buffer.h"
 #include "../ports/SkOSLibrary.h"
 
+namespace {
+
 typedef void *EGLDisplay;
 typedef unsigned int EGLBoolean;
 typedef void *EGLConfig;
@@ -25,6 +28,7 @@
 typedef void* EGLNativeWindowType;
 typedef void (*__eglMustCastToProperFunctionPointerType)(void);
 #define EGL_FALSE 0
+#define EGL_TRUE 1
 #define EGL_OPENGL_ES2_BIT 0x0004
 #define EGL_CONTEXT_CLIENT_VERSION 0x3098
 #define EGL_NO_SURFACE ((EGLSurface)0)
@@ -45,6 +49,8 @@
 #define EGL_NONE 0x3038
 #define EGL_WIDTH 0x3057
 #define EGL_HEIGHT 0x3056
+#define EGL_DRAW 0x3059
+#define EGL_READ 0x305A
 
 typedef EGLDisplay (*GetDisplayProc)(EGLNativeDisplayType display_id);
 typedef EGLBoolean (*InitializeProc)(EGLDisplay dpy, EGLint *major, EGLint *minor);
@@ -77,7 +83,6 @@
 static void* gLibrary = nullptr;
 static bool gfFunctionsLoadedSuccessfully = false;
 
-namespace {
 static void load_command_buffer_functions() {
     if (!gLibrary) {
         static constexpr const char* libName =
@@ -104,12 +109,11 @@
             gfSwapBuffers = (SwapBuffersProc)GetProcedureAddress(gLibrary, "eglSwapBuffers");
             gfGetProcAddress = (GetProcAddressProc)GetProcedureAddress(gLibrary, "eglGetProcAddress");
 
-            gfFunctionsLoadedSuccessfully = gfGetDisplay && gfInitialize && gfTerminate &&
-                                            gfChooseConfig && gfCreateWindowSurface &&
-                                            gfCreatePbufferSurface && gfDestroySurface &&
-                                            gfCreateContext && gfDestroyContext && gfMakeCurrent &&
-                                            gfSwapBuffers && gfGetProcAddress;
-
+            gfFunctionsLoadedSuccessfully =
+                    gfGetDisplay && gfInitialize && gfTerminate && gfChooseConfig &&
+                    gfCreateWindowSurface && gfCreatePbufferSurface && gfDestroySurface &&
+                    gfCreateContext && gfDestroyContext && gfMakeCurrent && gfSwapBuffers &&
+                    gfGetProcAddress;
         }
     }
 }
@@ -134,6 +138,80 @@
     return GrGLAssembleGLESInterface(gLibrary, command_buffer_get_gl_proc);
 }
 
+
+// The command buffer does not correctly implement eglGetCurrent. It always returns EGL_NO_<foo>.
+// So we implement them ourselves and hook eglMakeCurrent to store the current values in TLS.
+class TLSCurrentObjects {
+public:
+    static EGLDisplay CurrentDisplay() {
+        if (auto objects = Get()) {
+            return objects->fDisplay;
+        }
+        return EGL_NO_DISPLAY;
+    }
+
+    static EGLSurface CurrentSurface(EGLint readdraw) {
+        if (auto objects = Get()) {
+            switch (readdraw) {
+                case EGL_DRAW:
+                    return objects->fDrawSurface;
+                case EGL_READ:
+                    return objects->fReadSurface;
+                default:
+                    return EGL_NO_SURFACE;
+            }
+        }
+        return EGL_NO_SURFACE;
+    }
+
+    static EGLContext CurrentContext() {
+        if (auto objects = Get()) {
+            return objects->fContext;
+        }
+        return EGL_NO_CONTEXT;
+    }
+
+    static EGLBoolean MakeCurrent(EGLDisplay display, EGLSurface draw, EGLSurface read,
+                                  EGLContext ctx) {
+        if (gfFunctionsLoadedSuccessfully && EGL_TRUE == gfMakeCurrent(display, draw, read, ctx)) {
+            if (auto objects = Get()) {
+                objects->fDisplay = display;
+                objects->fDrawSurface = draw;
+                objects->fReadSurface = read;
+                objects->fContext = ctx;
+            }
+            return EGL_TRUE;
+        }
+        return EGL_FALSE;
+
+    }
+
+private:
+    EGLDisplay fDisplay = EGL_NO_DISPLAY;
+    EGLSurface fReadSurface = EGL_NO_SURFACE;
+    EGLSurface fDrawSurface = EGL_NO_SURFACE;
+    EGLContext fContext = EGL_NO_CONTEXT;
+
+    static TLSCurrentObjects* Get() {
+        return (TLSCurrentObjects*) SkTLS::Get(TLSCreate, TLSDelete);
+    }
+    static void* TLSCreate() { return new TLSCurrentObjects(); }
+    static void TLSDelete(void* objs) { delete (TLSCurrentObjects*)objs; }
+};
+
+std::function<void()> context_restorer() {
+    if (!gfFunctionsLoadedSuccessfully) {
+        return nullptr;
+    }
+    auto display = TLSCurrentObjects::CurrentDisplay();
+    auto dsurface = TLSCurrentObjects::CurrentSurface(EGL_DRAW);
+    auto rsurface = TLSCurrentObjects::CurrentSurface(EGL_READ);
+    auto context = TLSCurrentObjects::CurrentContext();
+    return [display, dsurface, rsurface, context] {
+        TLSCurrentObjects::MakeCurrent(display, dsurface, rsurface, context);
+    };
+}
+
 }  // anonymous namespace
 
 namespace sk_gpu_test {
@@ -204,7 +282,8 @@
         return;
     }
 
-    if (!gfMakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
+    SkScopeExit restorer(context_restorer());
+    if (!TLSCurrentObjects::MakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
         SkDebugf("Command Buffer: Could not make EGL context current.\n");
         this->destroyGLContext();
         return;
@@ -237,15 +316,19 @@
     if (EGL_NO_DISPLAY == fDisplay) {
         return;
     }
+    bool wasCurrent = false;
     if (EGL_NO_CONTEXT != fContext) {
+        wasCurrent = (TLSCurrentObjects::CurrentContext() == fContext);
         gfDestroyContext(fDisplay, fContext);
         fContext = EGL_NO_CONTEXT;
     }
-    // Call MakeCurrent after destroying the context, so that the EGL implementation knows that
-    // the context is not used anymore after it is released from being current.  This way
-    // command buffer does not need to abandon the context before destruction, and no
-    // client-side errors are printed.
-    gfMakeCurrent(fDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+    if (wasCurrent) {
+        // Call MakeCurrent after destroying the context, so that the EGL implementation knows that
+        // the context is not used anymore after it is released from being current.This way the
+        // command buffer does not need to abandon the context before destruction, and no
+        // client-side errors are printed.
+        TLSCurrentObjects::MakeCurrent(fDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+    }
 
     if (EGL_NO_SURFACE != fSurface) {
         gfDestroySurface(fDisplay, fSurface);
@@ -258,11 +341,18 @@
     if (!gfFunctionsLoadedSuccessfully) {
         return;
     }
-    if (!gfMakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
+    if (!TLSCurrentObjects::MakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
         SkDebugf("Command Buffer: Could not make EGL context current.\n");
     }
 }
 
+std::function<void()> CommandBufferGLTestContext::onPlatformGetAutoContextRestore() const {
+    if (!gfFunctionsLoadedSuccessfully || TLSCurrentObjects::CurrentContext() == fContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void CommandBufferGLTestContext::onPlatformSwapBuffers() const {
     if (!gfFunctionsLoadedSuccessfully) {
         return;
@@ -288,7 +378,7 @@
 }
 
 bool CommandBufferGLTestContext::makeCurrent() {
-    return gfMakeCurrent(fDisplay, fSurface, fSurface, fContext) != EGL_FALSE;
+    return TLSCurrentObjects::MakeCurrent(fDisplay, fSurface, fSurface, fContext) != EGL_FALSE;
 }
 
 int CommandBufferGLTestContext::getStencilBits() {
diff --git a/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.h b/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.h
index 7582f16..6a631be 100644
--- a/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.h
+++ b/tools/gpu/gl/command_buffer/GLTestContext_command_buffer.h
@@ -42,6 +42,8 @@
 
     void onPlatformMakeCurrent() const override;
 
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
+
     void onPlatformSwapBuffers() const override;
 
     GrGLFuncPtr onPlatformGetProcAddress(const char *name) const override;
diff --git a/tools/gpu/gl/debug/DebugGLTestContext.cpp b/tools/gpu/gl/debug/DebugGLTestContext.cpp
index f16692e..e28a3a7 100644
--- a/tools/gpu/gl/debug/DebugGLTestContext.cpp
+++ b/tools/gpu/gl/debug/DebugGLTestContext.cpp
@@ -1203,6 +1203,7 @@
 
 private:
     void onPlatformMakeCurrent() const override {}
+    std::function<void()> onPlatformGetAutoContextRestore() const override { return nullptr; }
     void onPlatformSwapBuffers() const override {}
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override { return nullptr; }
 };
diff --git a/tools/gpu/gl/egl/CreatePlatformGLTestContext_egl.cpp b/tools/gpu/gl/egl/CreatePlatformGLTestContext_egl.cpp
index 74cadfc..7fa88c0 100644
--- a/tools/gpu/gl/egl/CreatePlatformGLTestContext_egl.cpp
+++ b/tools/gpu/gl/egl/CreatePlatformGLTestContext_egl.cpp
@@ -39,6 +39,16 @@
     typedef sk_gpu_test::FenceSync INHERITED;
 };
 
+std::function<void()> context_restorer() {
+    auto display = eglGetCurrentDisplay();
+    auto dsurface = eglGetCurrentSurface(EGL_DRAW);
+    auto rsurface = eglGetCurrentSurface(EGL_READ);
+    auto context = eglGetCurrentContext();
+    return [display, dsurface, rsurface, context] {
+        eglMakeCurrent(display, dsurface, rsurface, context);
+    };
+}
+
 class EGLGLTestContext : public sk_gpu_test::GLTestContext {
 public:
     EGLGLTestContext(GrGLStandard forcedGpuAPI, EGLGLTestContext* shareContext);
@@ -53,6 +63,7 @@
     void destroyGLContext();
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override;
 
@@ -168,6 +179,7 @@
             continue;
         }
 
+        SkScopeExit restorer(context_restorer());
         if (!eglMakeCurrent(fDisplay, fSurface, fSurface, fContext)) {
             SkDebugf("eglMakeCurrent failed.  EGL Error: 0x%08x\n", eglGetError());
             this->destroyGLContext();
@@ -199,9 +211,11 @@
 
 void EGLGLTestContext::destroyGLContext() {
     if (fDisplay) {
-        eglMakeCurrent(fDisplay, 0, 0, 0);
-
         if (fContext) {
+            if (eglGetCurrentContext() == fContext) {
+                // This will ensure that the context is immediately deleted.
+                eglMakeCurrent(fDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+            }
             eglDestroyContext(fDisplay, fContext);
             fContext = EGL_NO_CONTEXT;
         }
@@ -284,6 +298,13 @@
     }
 }
 
+std::function<void()> EGLGLTestContext::onPlatformGetAutoContextRestore() const {
+    if (eglGetCurrentContext() == fContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void EGLGLTestContext::onPlatformSwapBuffers() const {
     if (!eglSwapBuffers(fDisplay, fSurface)) {
         SkDebugf("Could not complete eglSwapBuffers.\n");
diff --git a/tools/gpu/gl/glx/CreatePlatformGLTestContext_glx.cpp b/tools/gpu/gl/glx/CreatePlatformGLTestContext_glx.cpp
index 76b6d21..a6b2637 100644
--- a/tools/gpu/gl/glx/CreatePlatformGLTestContext_glx.cpp
+++ b/tools/gpu/gl/glx/CreatePlatformGLTestContext_glx.cpp
@@ -62,6 +62,7 @@
                                         GLXContext glxSharedContext);
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override;
 
@@ -90,6 +91,17 @@
     return ad->display();
 }
 
+std::function<void()> context_restorer() {
+    auto display = glXGetCurrentDisplay();
+    auto drawable = glXGetCurrentDrawable();
+    auto context = glXGetCurrentContext();
+    // On some systems calling glXMakeCurrent with a null display crashes.
+    if (!display) {
+        display = get_display();
+    }
+    return [display, drawable, context] { glXMakeCurrent(display, drawable, context); };
+}
+
 GLXGLTestContext::GLXGLTestContext(GrGLStandard forcedGpuAPI, GLXGLTestContext* shareContext)
     : fContext(nullptr)
     , fDisplay(nullptr)
@@ -214,6 +226,7 @@
         //SkDebugf("Direct GLX rendering context obtained.\n");
     }
 
+    SkScopeExit restorer(context_restorer());
     //SkDebugf("Making context current.\n");
     if (!glXMakeCurrent(fDisplay, fGlxPixmap, fContext)) {
       SkDebugf("Could not set the context.\n");
@@ -245,9 +258,11 @@
 
 void GLXGLTestContext::destroyGLContext() {
     if (fDisplay) {
-        glXMakeCurrent(fDisplay, 0, 0);
-
         if (fContext) {
+            if (glXGetCurrentContext() == fContext) {
+                // This will ensure that the context is immediately deleted.
+                glXMakeContextCurrent(fDisplay, None, None, nullptr);
+            }
             glXDestroyContext(fDisplay, fContext);
             fContext = nullptr;
         }
@@ -334,6 +349,13 @@
     }
 }
 
+std::function<void()> GLXGLTestContext::onPlatformGetAutoContextRestore() const {
+    if (glXGetCurrentContext() == fContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void GLXGLTestContext::onPlatformSwapBuffers() const {
     glXSwapBuffers(fDisplay, fGlxPixmap);
 }
diff --git a/tools/gpu/gl/iOS/CreatePlatformGLTestContext_iOS.mm b/tools/gpu/gl/iOS/CreatePlatformGLTestContext_iOS.mm
index e897e8c..65d2861 100644
--- a/tools/gpu/gl/iOS/CreatePlatformGLTestContext_iOS.mm
+++ b/tools/gpu/gl/iOS/CreatePlatformGLTestContext_iOS.mm
@@ -14,6 +14,11 @@
 
 namespace {
 
+std::function<void()> context_restorer() {
+    EAGLContext* context = [EAGLContext currentContext];
+    return [context] { [EAGLContext setCurrentContext:context]; };
+}
+
 class IOSGLTestContext : public sk_gpu_test::GLTestContext {
 public:
     IOSGLTestContext(IOSGLTestContext* shareContext);
@@ -23,6 +28,7 @@
     void destroyGLContext();
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override;
 
@@ -41,6 +47,7 @@
     } else {
         fEAGLContext = [[EAGLContext alloc] initWithAPI:kEAGLRenderingAPIOpenGLES2];
     }
+    SkScopeExit restorer(context_restorer());
     [EAGLContext setCurrentContext:fEAGLContext];
 
     sk_sp<const GrGLInterface> gl(GrGLCreateNativeInterface());
@@ -70,6 +77,7 @@
 void IOSGLTestContext::destroyGLContext() {
     if (fEAGLContext) {
         if ([EAGLContext currentContext] == fEAGLContext) {
+            // This will ensure that the context is immediately deleted.
             [EAGLContext setCurrentContext:nil];
         }
         fEAGLContext = nil;
@@ -86,6 +94,13 @@
     }
 }
 
+std::function<void()> IOSGLTestContext::onPlatformGetAutoContextRestore() const {
+    if ([EAGLContext currentContext] == fEAGLContext) {
+		return nullptr;
+	}
+    return context_restorer();
+}
+
 void IOSGLTestContext::onPlatformSwapBuffers() const { }
 
 GrGLFuncPtr IOSGLTestContext::onPlatformGetProcAddress(const char* procName) const {
diff --git a/tools/gpu/gl/mac/CreatePlatformGLTestContext_mac.cpp b/tools/gpu/gl/mac/CreatePlatformGLTestContext_mac.cpp
index a94f503..9f1c61e 100644
--- a/tools/gpu/gl/mac/CreatePlatformGLTestContext_mac.cpp
+++ b/tools/gpu/gl/mac/CreatePlatformGLTestContext_mac.cpp
@@ -14,6 +14,12 @@
 #include <dlfcn.h>
 
 namespace {
+
+std::function<void()> context_restorer() {
+    auto context = CGLGetCurrentContext();
+    return [context] { CGLSetCurrentContext(context); };
+}
+
 class MacGLTestContext : public sk_gpu_test::GLTestContext {
 public:
     MacGLTestContext(MacGLTestContext* shareContext);
@@ -23,6 +29,7 @@
     void destroyGLContext();
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override;
 
@@ -58,6 +65,7 @@
         return;
     }
 
+    SkScopeExit restorer(context_restorer());
     CGLSetCurrentContext(fContext);
 
     sk_sp<const GrGLInterface> gl(GrGLCreateNativeInterface());
@@ -86,6 +94,10 @@
 
 void MacGLTestContext::destroyGLContext() {
     if (fContext) {
+        if (CGLGetCurrentContext() == fContext) {
+            // This will ensure that the context is immediately deleted.
+            CGLSetCurrentContext(nullptr);
+        }
         CGLReleaseContext(fContext);
         fContext = nullptr;
     }
@@ -98,6 +110,13 @@
     CGLSetCurrentContext(fContext);
 }
 
+std::function<void()> MacGLTestContext::onPlatformGetAutoContextRestore() const {
+    if (CGLGetCurrentContext() == fContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void MacGLTestContext::onPlatformSwapBuffers() const {
     CGLFlushDrawable(fContext);
 }
diff --git a/tools/gpu/gl/null/NullGLTestContext.cpp b/tools/gpu/gl/null/NullGLTestContext.cpp
index 894de07..9e7279b 100644
--- a/tools/gpu/gl/null/NullGLTestContext.cpp
+++ b/tools/gpu/gl/null/NullGLTestContext.cpp
@@ -22,6 +22,7 @@
 
 private:
     void onPlatformMakeCurrent() const override {}
+    std::function<void()> onPlatformGetAutoContextRestore() const override { return nullptr; }
     void onPlatformSwapBuffers() const override {}
     GrGLFuncPtr onPlatformGetProcAddress(const char*) const override { return nullptr; }
 };
diff --git a/tools/gpu/gl/win/CreatePlatformGLTestContext_win.cpp b/tools/gpu/gl/win/CreatePlatformGLTestContext_win.cpp
index 0e97153..5fc355a 100644
--- a/tools/gpu/gl/win/CreatePlatformGLTestContext_win.cpp
+++ b/tools/gpu/gl/win/CreatePlatformGLTestContext_win.cpp
@@ -16,6 +16,12 @@
 
 namespace {
 
+std::function<void()> context_restorer() {
+    auto glrc = wglGetCurrentContext();
+    auto dc = wglGetCurrentDC();
+    return [glrc, dc] { wglMakeCurrent(dc, glrc); };
+}
+
 class WinGLTestContext : public sk_gpu_test::GLTestContext {
 public:
     WinGLTestContext(GrGLStandard forcedGpuAPI, WinGLTestContext* shareContext);
@@ -25,6 +31,7 @@
     void destroyGLContext();
 
     void onPlatformMakeCurrent() const override;
+    std::function<void()> onPlatformGetAutoContextRestore() const override;
     void onPlatformSwapBuffers() const override;
     GrGLFuncPtr onPlatformGetProcAddress(const char* name) const override;
 
@@ -113,6 +120,7 @@
         glrc = fPbufferContext->getGLRC();
     }
 
+    SkScopeExit restorer(context_restorer());
     if (!(wglMakeCurrent(dc, glrc))) {
         SkDebugf("Could not set the context.\n");
         this->destroyGLContext();
@@ -142,6 +150,7 @@
 void WinGLTestContext::destroyGLContext() {
     SkSafeSetNull(fPbufferContext);
     if (fGlRenderContext) {
+        // This deletes the context immediately even if it is current.
         wglDeleteContext(fGlRenderContext);
         fGlRenderContext = 0;
     }
@@ -172,6 +181,13 @@
     }
 }
 
+std::function<void()> WinGLTestContext::onPlatformGetAutoContextRestore() const {
+    if (wglGetCurrentContext() == fGlRenderContext) {
+        return nullptr;
+    }
+    return context_restorer();
+}
+
 void WinGLTestContext::onPlatformSwapBuffers() const {
     HDC dc;
 
diff --git a/tools/gpu/mock/MockTestContext.cpp b/tools/gpu/mock/MockTestContext.cpp
index 68941ad..5cc5529 100644
--- a/tools/gpu/mock/MockTestContext.cpp
+++ b/tools/gpu/mock/MockTestContext.cpp
@@ -34,6 +34,7 @@
 protected:
     void teardown() override {}
     void onPlatformMakeCurrent() const override {}
+    std::function<void()> onPlatformGetAutoContextRestore() const override { return nullptr; }
     void onPlatformSwapBuffers() const override {}
 
 private:
diff --git a/tools/gpu/mtl/MtlTestContext.mm b/tools/gpu/mtl/MtlTestContext.mm
index 4014e2b..38755b2 100644
--- a/tools/gpu/mtl/MtlTestContext.mm
+++ b/tools/gpu/mtl/MtlTestContext.mm
@@ -144,6 +144,7 @@
     }
 
     void onPlatformMakeCurrent() const override {}
+    std::function<void()> onPlatformGetAutoContextRestore() const override { return nullptr; }
     void onPlatformSwapBuffers() const override {}
 
     id<MTLDevice> fDevice;
diff --git a/tools/gpu/vk/VkTestContext.cpp b/tools/gpu/vk/VkTestContext.cpp
index daaa55b..fa40c73 100644
--- a/tools/gpu/vk/VkTestContext.cpp
+++ b/tools/gpu/vk/VkTestContext.cpp
@@ -154,6 +154,7 @@
     }
 
     void onPlatformMakeCurrent() const override {}
+    std::function<void()> onPlatformGetAutoContextRestore() const override  { return nullptr; }
     void onPlatformSwapBuffers() const override {}
 
     typedef sk_gpu_test::VkTestContext INHERITED;