From 60ca1a0d50011f43f5ba97f2e1c5556752957f42 Mon Sep 17 00:00:00 2001
From: Antoni Sawicki <tenox@google.com>
Date: Sun, 26 Apr 2020 00:23:31 -0700
Subject: [PATCH] lame attempt at restarting cancelled context

---
 wrp.go | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/wrp.go b/wrp.go
index 289fa4a..149b4a1 100644
--- a/wrp.go
+++ b/wrp.go
@@ -318,11 +318,17 @@ func (w wrpReq) capture() {
 	} else {
 		chromedp.Run(ctx, emulation.SetDeviceMetricsOverride(int64(float64(w.width)/w.scale), int64(float64(w.height)/w.scale), w.scale, false))
 	}
+	// Capture screenshot...
 	err = chromedp.Run(ctx, chromedp.CaptureScreenshot(&pngcap))
 	if err != nil {
-		// TODO: process context cancelled here
-		log.Printf("%s Failed to capture screenshot: %s\n", w.req.RemoteAddr, err)
-		fmt.Fprintf(w.out, "<BR>Unable to capture screenshot:<BR>%s<BR>\n", err)
+		if err.Error() == "context canceled" {
+			log.Printf("%s Contex cancelled, try again", w.req.RemoteAddr)
+			fmt.Fprintf(w.out, "<BR>%s<BR> -- restarting, try again", err)
+			ctx, cancel = chromedp.NewContext(context.Background())
+		} else {
+			log.Printf("%s Failed to capture screenshot: %s\n", w.req.RemoteAddr, err)
+			fmt.Fprintf(w.out, "<BR>Unable to capture screenshot:<BR>%s<BR>\n", err)
+		}
 		return
 	}
 	seq := rand.Intn(9999)