[host] simplify capture logic and fix re-init bug

host/Service.cpp

@@ -170,66 +170,7 @@ void Service::DeInitialize()
   m_initialized = false;
 }
 
-bool Service::Process()
+bool Service::ReInit(volatile char * flags)
-{
-  if (!m_initialized)
-    return false;
-
-  volatile char * flags = (volatile char *)&(m_shmHeader->flags);
-
-  // check if the client has flagged a restart
-  if (*flags & KVMFR_HEADER_FLAG_RESTART)
-  {
-    DEBUG_INFO("Restart Requested");
-    if (!m_capture->ReInitialize())
-    {
-      DEBUG_ERROR("ReInitialize Failed");
-      return false;
-    }
-
-    if (m_capture->GetMaxFrameSize() > m_frameSize)
-    {
-      DEBUG_ERROR("Maximum frame size of %zd bytes exceeds maximum space available", m_capture->GetMaxFrameSize());
-      return false;
-    }
-
-    INTERLOCKED_AND8(flags, ~(KVMFR_HEADER_FLAG_RESTART));
-  }
-
-  unsigned int status;
-  bool ok         = false;
-  bool repeat     = false;
-
-  for(int i = 0; i < 2; ++i)
-  {
-    status = m_capture->Capture();
-    if (status & GRAB_STATUS_OK)
-    {
-      ok = true;
-      break;
-    }
-
-    if (status & GRAB_STATUS_TIMEOUT)
-    {
-      if (m_haveFrame)
-      {
-        ok     = true;
-        repeat = true;
-        break;
-      }
-
-      // timeouts should not count towards a failure to capture
-      --i;
-      continue;
-    }
-
-    if (status & GRAB_STATUS_ERROR)
-    {
-      DEBUG_ERROR("Capture failed, retrying");
-      continue;
-    }
-
-    if (status & GRAB_STATUS_REINIT)
 {
   DEBUG_INFO("ReInitialize Requested");
 
@@ -257,32 +198,74 @@ bool Service::Process()
   }
 
   INTERLOCKED_AND8(flags, ~KVMFR_HEADER_FLAG_PAUSED);
+  return true;
+}
+
+ProcessStatus Service::Process()
+{
+  if (!m_initialized)
+    return PROCESS_STATUS_ERROR;
+
+  volatile char * flags = (volatile char *)&(m_shmHeader->flags);
+
+  // check if the client has flagged a restart
+  if (*flags & KVMFR_HEADER_FLAG_RESTART)
+  {
+    DEBUG_INFO("Restart Requested");
+    if (!m_capture->ReInitialize())
+    {
+      DEBUG_ERROR("ReInitialize Failed");
+      return PROCESS_STATUS_ERROR;
+    }
+
+    if (m_capture->GetMaxFrameSize() > m_frameSize)
+    {
+      DEBUG_ERROR("Maximum frame size of %zd bytes exceeds maximum space available", m_capture->GetMaxFrameSize());
+      return PROCESS_STATUS_ERROR;
+    }
+
+    INTERLOCKED_AND8(flags, ~(KVMFR_HEADER_FLAG_RESTART));
+  }
+
+  unsigned int status;
+  bool notify = false;
+
+  status = m_capture->Capture();
+  if (status & GRAB_STATUS_ERROR)
+  {
+    DEBUG_WARN("Capture error, retrying");
+    return PROCESS_STATUS_RETRY;
+  }
+
+  if (status & GRAB_STATUS_TIMEOUT)
+  {
+    // timeouts should not count towards a failure to capture
+    if (!m_haveFrame)
+      return PROCESS_STATUS_OK;
+
+    notify = true;
+  }
+
+  if (status & GRAB_STATUS_REINIT)
+  {
+    if (!ReInit(flags))
+      return PROCESS_STATUS_ERROR;
 
     // re-init request should not count towards a failure to capture
-      --i;
+    return PROCESS_STATUS_OK;
-      continue;
   }
 
-    DEBUG_ERROR("Capture interface returned an unexpected result");
+  if ((status & (GRAB_STATUS_OK | GRAB_STATUS_TIMEOUT)) == 0)
-    return false;
-  }
-
-
-  if (!ok)
   {
-    DEBUG_ERROR("Capture retry count exceeded");
+    DEBUG_ERROR("Capture interface returned an unexpected result");
-    return false;
+    return PROCESS_STATUS_ERROR;
   }
 
   if (status & GRAB_STATUS_CURSOR)
     SetEvent(m_cursorEvent);
 
-  if (status & GRAB_STATUS_FRAME)
-  {
   volatile KVMFRFrame * fi = &(m_shmHeader->frame);
- 
+  if (status & GRAB_STATUS_FRAME)
-    // only update the header if the frame is new
-    if (!repeat)
   { 
     FrameInfo frame  = { 0 };
     frame.buffer     = m_frame[m_frameIndex];
@@ -291,8 +274,17 @@ bool Service::Process()
     GrabStatus result = m_capture->GetFrame(frame);
     if (result != GRAB_STATUS_OK)
     {
+      if (result == GRAB_STATUS_REINIT)
+      {
+        if (!ReInit(flags))
+          return PROCESS_STATUS_ERROR;
+
+        // re-init request should not count towards a failure to capture
+        return PROCESS_STATUS_OK;
+      }
+
       DEBUG_INFO("GetFrame failed");
-        return false;
+      return PROCESS_STATUS_ERROR;
     }
 
     /* don't touch the frame information until the client is done with it */
@@ -316,8 +308,10 @@ bool Service::Process()
 
     // remember that we have a valid frame
     m_haveFrame = true;
+    notify = true;
   }
-    else
+
+  if (notify)
   {
     /* don't touch the frame inforamtion until the client is done with it */
     while (fi->flags & KVMFR_FRAME_FLAG_UPDATE)
@@ -325,15 +319,13 @@ bool Service::Process()
       if (*flags & KVMFR_HEADER_FLAG_RESTART)
         break;
     }
-    }
-
     // signal a frame update
     fi->flags |= KVMFR_FRAME_FLAG_UPDATE;
   }
 
   // update the flags
   INTERLOCKED_AND8(flags, KVMFR_HEADER_FLAG_RESTART);
-  return true;
+  return PROCESS_STATUS_OK;
 }
 
 DWORD Service::CursorThread()

host/Service.h

@@ -28,6 +28,13 @@ Place, Suite 330, Boston, MA 02111-1307 USA
 
 #define MAX_FRAMES 2
 
+enum ProcessStatus
+{
+  PROCESS_STATUS_OK,
+  PROCESS_STATUS_RETRY,
+  PROCESS_STATUS_ERROR
+};
+
 class Service
 {
 public:
@@ -40,7 +47,7 @@ public:
 
   bool Initialize(ICapture * captureDevice);
   void DeInitialize();
-  bool Process();
+  ProcessStatus Process();
 
 private:
   bool InitPointers();
@@ -52,6 +59,8 @@ private:
   Service();
   ~Service();
 
+  bool ReInit(volatile char * flags);
+
   bool       m_initialized;
   bool       m_running;
   DWORD      m_consoleSessionID;

host/main.cpp

@@ -128,10 +128,28 @@ int run(struct StartupArgs & args)
   if (!svc->Initialize(captureDevice))
     return -1;
 
-  while (true)
+  int retry = 0;
+  bool running = true;
+  while (running)
   {
-    if (!svc->Process())
+    switch (svc->Process())
+    {
+      case PROCESS_STATUS_OK:
+        retry = 0;
         break;
+
+      case PROCESS_STATUS_RETRY:
+        if (retry++ == 3)
+        {
+          fprintf(stderr, "Too many consecutive retries, aborting");
+          running = false;
+        }
+        break;
+
+      case PROCESS_STATUS_ERROR:
+        fprintf(stderr, "Capture process returned error");
+        running = false;
+    }
   }
 
   svc->DeInitialize();