From 9ce3d2dbf079451bc44339f61f442434c36aee2f Mon Sep 17 00:00:00 2001
From: Antoine Martin <totaam@xpra.org>
Date: Thu, 27 Apr 2017 16:33:44 +0000
Subject: [PATCH] #389: * try to handle transient failures and re-init the
 context * return an exit code from the test capture tool

git-svn-id: https://xpra.org/svn/Xpra/trunk@15727 3bb7dfac-3a0b-4e04-842a-767bc560f471
---
 src/xpra/codecs/codec_constants.py       |  3 ++
 src/xpra/codecs/nvfbc/capture.py         | 36 +++++++++++++-----------
 src/xpra/codecs/nvfbc/fbc_capture.pyx    | 12 ++++----
 src/xpra/platform/win32/shadow_server.py | 18 +++++++++++-
 4 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/src/xpra/codecs/codec_constants.py b/src/xpra/codecs/codec_constants.py
index 5e80faeadc..f907a21dbc 100644
--- a/src/xpra/codecs/codec_constants.py
+++ b/src/xpra/codecs/codec_constants.py
@@ -42,6 +42,9 @@ def get_subsampling_divs(pixel_format):
 class TransientCodecException(Exception):
     pass
 
+class CodecStateException(Exception):
+    pass
+
 
 class _codec_spec(object):
 
diff --git a/src/xpra/codecs/nvfbc/capture.py b/src/xpra/codecs/nvfbc/capture.py
index 3003b60147..20e114cfe2 100644
--- a/src/xpra/codecs/nvfbc/capture.py
+++ b/src/xpra/codecs/nvfbc/capture.py
@@ -34,23 +34,25 @@ def main():
             log("Capture()", exc_info=True)
             log.error("Error: failed to create test capture instance:")
             log.error(" %s", e)
-        else:
-            image = c.get_image()
-            assert image
-            from PIL import Image
-            w = image.get_width()
-            h = image.get_height()
-            pixels = image.get_pixels()
-            stride = image.get_rowstride()
-            rgb_format = image.get_pixel_format()
-            try:
-                img = Image.frombuffer("RGB", (w, h), pixels, "raw", rgb_format, stride, 1)
-                filename = os.path.join(get_download_dir(), "screenshot-%s-%i.png" % (rgb_format, time.time()))
-                img.save(filename, "png")
-                log.info("screenshot saved to %s", filename)
-            except Exception as e:
-                log.warn("not saved %s: %s", rgb_format, e)
+            return 1
+        image = c.get_image()
+        assert image
+        from PIL import Image
+        w = image.get_width()
+        h = image.get_height()
+        pixels = image.get_pixels()
+        stride = image.get_rowstride()
+        rgb_format = image.get_pixel_format()
+        try:
+            img = Image.frombuffer("RGB", (w, h), pixels, "raw", rgb_format, stride, 1)
+            filename = os.path.join(get_download_dir(), "screenshot-%s-%i.png" % (rgb_format, time.time()))
+            img.save(filename, "png")
+            log.info("screenshot saved to %s", filename)
+            return 0
+        except Exception as e:
+            log.warn("not saved %s: %s", rgb_format, e)
+            return 1
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/src/xpra/codecs/nvfbc/fbc_capture.pyx b/src/xpra/codecs/nvfbc/fbc_capture.pyx
index 896f020449..67dcabf7b6 100644
--- a/src/xpra/codecs/nvfbc/fbc_capture.pyx
+++ b/src/xpra/codecs/nvfbc/fbc_capture.pyx
@@ -9,7 +9,7 @@ import sys
 
 from xpra.os_util import WIN32
 from xpra.codecs.image_wrapper import ImageWrapper
-from xpra.codecs.codec_constants import TransientCodecException
+from xpra.codecs.codec_constants import TransientCodecException, CodecStateException
 from xpra.codecs.nv_util import get_nvidia_module_version, get_cards
 
 from xpra.log import Logger
@@ -474,12 +474,12 @@ cdef class NvFBC_Capture:
         cdef NVFBCRESULT res
         with nogil:
             res = self.context.NvFBCToSysGrabFrame(&grab)
-        if res!=0 and grab_info.bMustRecreate:
-            pass    #TODO!
         if res!=0 and grab_info.dwDriverInternalError:
-            pass
-        if res==NVFBC_ERROR_INVALIDATED_SESSION:
-            pass    #retry!
+            raise CodecStateException("NvFBC driver internal error")
+        if res==NVFBC_ERROR_DYNAMIC_DISABLE:
+            raise CodecStateException("NvFBC capture has been disabled")
+        if (res!=0 and grab_info.bMustRecreate) or res==NVFBC_ERROR_INVALIDATED_SESSION:
+            raise TransientCodecException("NvFBC context invalidated")
         log("NvFBCToSysGrabFrame(%#x)=%i", <uintptr_t> &grab, res)
         raiseNvFBC(res, "NvFBCToSysGrabFrame")
         info = {
diff --git a/src/xpra/platform/win32/shadow_server.py b/src/xpra/platform/win32/shadow_server.py
index 0a122c5243..a5441f851b 100644
--- a/src/xpra/platform/win32/shadow_server.py
+++ b/src/xpra/platform/win32/shadow_server.py
@@ -18,6 +18,7 @@
 
 from xpra.util import XPRA_APP_ID
 from xpra.scripts.config import InitException
+from xpra.codecs.codec_constants import CodecStateException, TransientCodecException
 from xpra.server.gtk_server_base import GTKServerBase
 from xpra.server.shadow.gtk_shadow_server_base import GTKShadowServerBase
 from xpra.server.shadow.root_window_model import RootWindowModel
@@ -94,6 +95,9 @@ def init_capture(self):
 
     def cleanup(self):
         RootWindowModel.cleanup(self)
+        self.cleanup_capture()
+
+    def cleanup_capture(self):
         c = self.capture
         if c:
             self.capture = None
@@ -225,7 +229,19 @@ def get_root_window_size(self):
         return w, h
 
     def get_image(self, x, y, width, height, logger=None):
-        return self.capture.get_image(x, y, width, height)
+        try:
+            return self.capture.get_image(x, y, width, height)
+        except CodecStateException as e:
+            #maybe we should exit here?
+            log.warn("Warning: %s", e)
+            self.cleanup_capture()
+            self.init_capture()
+            return None
+        except TransientCodecException as e:
+            log.warn("Warning: %s", e)
+            self.cleanup_capture()
+            self.init_capture()
+            return None
 
     def take_screenshot(self):
         return self.capture.take_screenshot()