diff --git a/services/java/com/android/server/WifiService.java b/services/java/com/android/server/WifiService.java
index 35b250e6ba44b..bfbd76184734e 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -232,18 +232,6 @@ public class WifiService extends IWifiManager.Stub {
         PowerManager powerManager = (PowerManager)mContext.getSystemService(Context.POWER_SERVICE);
         sWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, WAKELOCK_TAG);
         sDriverStopWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, WAKELOCK_TAG);
-        mWifiStateTracker.setReleaseWakeLockCallback(
-                new Runnable() {
-                    public void run() {
-                        mWifiHandler.removeMessages(MESSAGE_RELEASE_WAKELOCK);
-                        synchronized (sDriverStopWakeLock) {
-                            if (sDriverStopWakeLock.isHeld()) {
-                                sDriverStopWakeLock.release();
-                            }
-                        }
-                    }
-                }
-        );
 
         mContext.registerReceiver(
                 new BroadcastReceiver() {
@@ -1779,20 +1767,16 @@ public class WifiService extends IWifiManager.Stub {
                     sendEnableMessage(true, false, mLastEnableUid);
                     sWakeLock.acquire();
                     sendStartMessage(strongestLockMode == WifiManager.WIFI_MODE_SCAN_ONLY);
-                } else {
+                } else if (!mWifiStateTracker.isDriverStopped()) {
                     int wakeLockTimeout =
                             Settings.Secure.getInt(
                                     mContext.getContentResolver(),
                                     Settings.Secure.WIFI_MOBILE_DATA_TRANSITION_WAKELOCK_TIMEOUT_MS,
                                     DEFAULT_WAKELOCK_TIMEOUT);
                     /*
-                     * The following wakelock is held in order to ensure
-                     * that the connectivity manager has time to fail over
-                     * to the mobile data network. The connectivity manager
-                     * releases it once mobile data connectivity has been
-                     * established. If connectivity cannot be established,
-                     * the wakelock is released after wakeLockTimeout
-                     * milliseconds have elapsed.
+                     * We are assuming that ConnectivityService can make
+                     * a transition to cellular data within wakeLockTimeout time.
+                     * The wakelock is released by the delayed message.
                      */
                     sDriverStopWakeLock.acquire();
                     mWifiHandler.sendEmptyMessage(MESSAGE_STOP_WIFI);
@@ -1886,11 +1870,7 @@ public class WifiService extends IWifiManager.Stub {
                     break;
 
                 case MESSAGE_RELEASE_WAKELOCK:
-                    synchronized (sDriverStopWakeLock) {
-                        if (sDriverStopWakeLock.isHeld()) {
-                            sDriverStopWakeLock.release();
-                        }
-                    }
+                    sDriverStopWakeLock.release();
                     break;
 
                 case MESSAGE_START_ACCESS_POINT:
diff --git a/wifi/java/android/net/wifi/WifiStateTracker.java b/wifi/java/android/net/wifi/WifiStateTracker.java
index 80b8aedfa0d03..4f84aabe9d37a 100644
--- a/wifi/java/android/net/wifi/WifiStateTracker.java
+++ b/wifi/java/android/net/wifi/WifiStateTracker.java
@@ -313,8 +313,6 @@ public class WifiStateTracker extends NetworkStateTracker {
     private String mInterfaceName;
     private static String LS = System.getProperty("line.separator");
 
-    private Runnable mReleaseWakeLockCallback;
-
     private static String[] sDnsPropNames;
 
     /**
@@ -615,7 +613,15 @@ public class WifiStateTracker extends NetworkStateTracker {
         }
     }
 
-    private synchronized boolean isDriverStopped() {
+    /**
+     * TODO: mRunState is not synchronized in some places
+     * address this as part of re-architect.
+     *
+     * TODO: We are exposing an additional public synchronized call
+     * for a wakelock optimization in WifiService. Remove it
+     * when we handle the wakelock in ConnectivityService.
+     */
+    public synchronized boolean isDriverStopped() {
         return mRunState == RUN_STATE_STOPPED || mRunState == RUN_STATE_STOPPING;
     }
 
@@ -674,15 +680,20 @@ public class WifiStateTracker extends NetworkStateTracker {
         }
     }
 
+    /**
+     * We release the wakelock in WifiService
+     * using a timer.
+     *
+     * TODO:
+     * Releasing wakelock using both timer and
+     * a call from ConnectivityService requires
+     * a rethink. We had problems where WifiService
+     * could keep a wakelock forever if we delete
+     * messages in the asynchronous call
+     * from ConnectivityService
+     */
     @Override
     public void releaseWakeLock() {
-        if (mReleaseWakeLockCallback != null) {
-            mReleaseWakeLockCallback.run();
-        }
-    }
-    
-    public void setReleaseWakeLockCallback(Runnable callback) {
-        mReleaseWakeLockCallback = callback;
     }
 
     /**
@@ -1481,20 +1492,6 @@ public class WifiStateTracker extends NetworkStateTracker {
             } else {
                 return disconnect();
             }
-        } else {
-            /*
-             * The "driver-stop" wake lock normally is released from the
-             * connectivity manager after the mobile data connection has
-             * been established, or after a timeout period, if that never
-             * happens. Because WifiService.updateWifiState() can get called
-             * multiple times, we can end up acquiring the wake lock and calling
-             * disconnectAndStop() even when a disconnect or stop operation
-             * is already in progress. In that case, we want to ignore the
-             * disconnectAndStop request and release the (ref-counted) wake
-             * lock, so that eventually, when the mobile data connection is
-             * established, the ref count will drop to zero.
-             */
-            releaseWakeLock();
         }
         return true;
     }