From 05ecddd14016617011d9d40dfa73b5f622d36fd1 Mon Sep 17 00:00:00 2001
From: Nikhil Marathe <nsm.nikhil@gmail.com>
Date: Mon, 9 Mar 2015 15:40:04 -0700
Subject: [PATCH] Bug 1080752 - Hold wakelock when attempting to connect to
 push server. r=dougt, a=jocheng

---
 dom/push/PushService.jsm | 53 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/dom/push/PushService.jsm b/dom/push/PushService.jsm
index 70e7a525b64d..952fa7576de2 100644
--- a/dom/push/PushService.jsm
+++ b/dom/push/PushService.jsm
@@ -33,6 +33,10 @@ XPCOMUtils.defineLazyServiceGetter(this, "gDNSService",
 XPCOMUtils.defineLazyModuleGetter(this, "AlarmService",
                                   "resource://gre/modules/AlarmService.jsm");
 
+XPCOMUtils.defineLazyServiceGetter(this, "gPowerManagerService",
+                                   "@mozilla.org/power/powermanagerservice;1",
+                                   "nsIPowerManagerService");
+
 var threadManager = Cc["@mozilla.org/thread-manager;1"].getService(Ci.nsIThreadManager);
 
 this.EXPORTED_SYMBOLS = ["PushService"];
@@ -832,8 +836,18 @@ this.PushService = {
     debug("serverURL: " + uri.spec);
     this._wsListener = new PushWebSocketListener(this);
     this._ws.protocol = "push-notification";
-    this._ws.asyncOpen(uri, serverURL, this._wsListener, null);
-    this._currentState = STATE_WAITING_FOR_WS_START;
+
+    try {
+      // Grab a wakelock before we open the socket to ensure we don't go to sleep
+      // before connection the is opened.
+      this._ws.asyncOpen(uri, serverURL, this._wsListener, null);
+      this._acquireWakeLock();
+      this._currentState = STATE_WAITING_FOR_WS_START;
+    } catch(e) {
+      debug("Error opening websocket. asyncOpen failed!");
+      this._shutdownWS();
+      this._reconnectAfterBackoff();
+    }
   },
 
   _startListeningIfChannelsPresent: function() {
@@ -956,6 +970,38 @@ this.PushService = {
     }
   },
 
+  _acquireWakeLock: function() {
+    if (!this._socketWakeLock) {
+      debug("Acquiring Socket Wakelock");
+      this._socketWakeLock = gPowerManagerService.newWakeLock("cpu");
+    }
+    if (!this._socketWakeLockTimer) {
+      debug("Creating Socket WakeLock Timer");
+      this._socketWakeLockTimer = Cc["@mozilla.org/timer;1"].createInstance(Ci.nsITimer);
+    }
+
+    debug("Setting Socket WakeLock Timer");
+    this._socketWakeLockTimer
+      .initWithCallback(this._releaseWakeLock.bind(this),
+                        // Allow the same time for socket setup as we do for
+                        // requests after the setup. Fudge it a bit since
+                        // timers can be a little off and we don't want to go
+                        // to sleep just as the socket connected.
+                        this._requestTimeout + 1000,
+                        Ci.nsITimer.ONE_SHOT);
+  },
+
+  _releaseWakeLock: function() {
+    debug("Releasing Socket WakeLock");
+    if (this._socketWakeLockTimer) {
+      this._socketWakeLockTimer.cancel();
+    }
+    if (this._socketWakeLock) {
+      this._socketWakeLock.unlock();
+      this._socketWakeLock = null;
+    }
+  },
+
   /**
    * Protocol handler invoked by server message.
    */
@@ -1471,6 +1517,8 @@ this.PushService = {
   // begin Push protocol handshake
   _wsOnStart: function(context) {
     debug("wsOnStart()");
+    this._releaseWakeLock();
+
     if (this._currentState != STATE_WAITING_FOR_WS_START) {
       debug("NOT in STATE_WAITING_FOR_WS_START. Current state " +
             this._currentState + ". Skipping");
@@ -1527,6 +1575,7 @@ this.PushService = {
    */
   _wsOnStop: function(context, statusCode) {
     debug("wsOnStop()");
+    this._releaseWakeLock();
 
     if (statusCode != Cr.NS_OK &&
         !(statusCode == Cr.NS_BASE_STREAM_CLOSED && this._willBeWokenUpByUDP)) {
-- 
2.11.4.GIT