INetworkStatsSession with lifecycle for caching.

Users outside system_server now explicitly communicate their
lifecycle, which keeps a strong-reference chain to any fully loaded
NetworkStatsCollection histories.

Bug: 6236498
Change-Id: I8e22739b6e89a626b676967a736d7117fd000778
diff --git a/Android.mk b/Android.mk
index d27dbab..9c51fc6 100644
--- a/Android.mk
+++ b/Android.mk
@@ -118,6 +118,7 @@
 	core/java/android/net/INetworkPolicyListener.aidl \
 	core/java/android/net/INetworkPolicyManager.aidl \
 	core/java/android/net/INetworkStatsService.aidl \
+	core/java/android/net/INetworkStatsSession.aidl \
 	core/java/android/net/nsd/INsdManager.aidl \
 	core/java/android/nfc/INdefPushCallback.aidl \
 	core/java/android/nfc/INfcAdapter.aidl \
diff --git a/core/java/android/net/INetworkStatsService.aidl b/core/java/android/net/INetworkStatsService.aidl
index 0e883cf..b4f6367 100644
--- a/core/java/android/net/INetworkStatsService.aidl
+++ b/core/java/android/net/INetworkStatsService.aidl
@@ -16,6 +16,7 @@
 
 package android.net;
 
+import android.net.INetworkStatsSession;
 import android.net.NetworkStats;
 import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
@@ -23,15 +24,11 @@
 /** {@hide} */
 interface INetworkStatsService {
 
-    /** Return historical network layer stats for traffic that matches template. */
-    NetworkStatsHistory getHistoryForNetwork(in NetworkTemplate template, int fields);
-    /** Return historical network layer stats for specific UID traffic that matches template. */
-    NetworkStatsHistory getHistoryForUid(in NetworkTemplate template, int uid, int set, int tag, int fields);
+    /** Start a statistics query session. */
+    INetworkStatsSession openSession();
 
-    /** Return network layer usage summary for traffic that matches template. */
-    NetworkStats getSummaryForNetwork(in NetworkTemplate template, long start, long end);
-    /** Return network layer usage summary per UID for traffic that matches template. */
-    NetworkStats getSummaryForAllUid(in NetworkTemplate template, long start, long end, boolean includeTags);
+    /** Return network layer usage total for traffic that matches template. */
+    long getNetworkTotalBytes(in NetworkTemplate template, long start, long end);
 
     /** Return data layer snapshot of UID network usage. */
     NetworkStats getDataLayerSnapshotForUid(int uid);
diff --git a/core/java/android/net/INetworkStatsSession.aidl b/core/java/android/net/INetworkStatsSession.aidl
new file mode 100644
index 0000000..1596fa2
--- /dev/null
+++ b/core/java/android/net/INetworkStatsSession.aidl
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2012 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net;
+
+import android.net.NetworkStats;
+import android.net.NetworkStatsHistory;
+import android.net.NetworkTemplate;
+
+/** {@hide} */
+interface INetworkStatsSession {
+
+    /** Return network layer usage summary for traffic that matches template. */
+    NetworkStats getSummaryForNetwork(in NetworkTemplate template, long start, long end);
+    /** Return historical network layer stats for traffic that matches template. */
+    NetworkStatsHistory getHistoryForNetwork(in NetworkTemplate template, int fields);
+
+    /** Return network layer usage summary per UID for traffic that matches template. */
+    NetworkStats getSummaryForAllUid(in NetworkTemplate template, long start, long end, boolean includeTags);
+    /** Return historical network layer stats for specific UID traffic that matches template. */
+    NetworkStatsHistory getHistoryForUid(in NetworkTemplate template, int uid, int set, int tag, int fields);
+
+    void close();
+
+}
diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index 973fac1..ee3e165 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -238,6 +238,19 @@
         }
     }
 
+    /** {@hide} */
+    public static void closeQuietly(INetworkStatsSession session) {
+        // TODO: move to NetworkStatsService once it exists
+        if (session != null) {
+            try {
+                session.close();
+            } catch (RuntimeException rethrown) {
+                throw rethrown;
+            } catch (Exception ignored) {
+            }
+        }
+    }
+
     /**
      * Get the total number of packets transmitted through the mobile interface.
      *
diff --git a/services/java/com/android/server/net/NetworkPolicyManagerService.java b/services/java/com/android/server/net/NetworkPolicyManagerService.java
index b0657a6..77addd6 100644
--- a/services/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -1740,7 +1740,7 @@
 
     private long getTotalBytes(NetworkTemplate template, long start, long end) {
         try {
-            return mNetworkStats.getSummaryForNetwork(template, start, end).getTotalBytes();
+            return mNetworkStats.getNetworkTotalBytes(template, start, end);
         } catch (RuntimeException e) {
             Slog.w(TAG, "problem reading network stats: " + e);
             return 0;
diff --git a/services/java/com/android/server/net/NetworkStatsRecorder.java b/services/java/com/android/server/net/NetworkStatsRecorder.java
index 290bd2c..540f606 100644
--- a/services/java/com/android/server/net/NetworkStatsRecorder.java
+++ b/services/java/com/android/server/net/NetworkStatsRecorder.java
@@ -221,6 +221,11 @@
         if (mLastSnapshot != null) {
             mLastSnapshot = mLastSnapshot.withoutUid(uid);
         }
+
+        final NetworkStatsCollection complete = mComplete != null ? mComplete.get() : null;
+        if (complete != null) {
+            complete.removeUid(uid);
+        }
     }
 
     /**
diff --git a/services/java/com/android/server/net/NetworkStatsService.java b/services/java/com/android/server/net/NetworkStatsService.java
index b847673..f7ba329 100644
--- a/services/java/com/android/server/net/NetworkStatsService.java
+++ b/services/java/com/android/server/net/NetworkStatsService.java
@@ -70,6 +70,7 @@
 import android.net.IConnectivityManager;
 import android.net.INetworkManagementEventObserver;
 import android.net.INetworkStatsService;
+import android.net.INetworkStatsSession;
 import android.net.LinkProperties;
 import android.net.NetworkIdentity;
 import android.net.NetworkInfo;
@@ -412,40 +413,75 @@
     }
 
     @Override
-    public NetworkStatsHistory getHistoryForNetwork(NetworkTemplate template, int fields) {
-        return mDevStatsCached.getHistory(template, UID_ALL, SET_ALL, TAG_NONE, fields);
+    public INetworkStatsSession openSession() {
+        mContext.enforceCallingOrSelfPermission(READ_NETWORK_USAGE_HISTORY, TAG);
+
+        // return an IBinder which holds strong references to any loaded stats
+        // for its lifetime; when caller closes only weak references remain.
+
+        return new INetworkStatsSession.Stub() {
+            private NetworkStatsCollection mUidComplete;
+            private NetworkStatsCollection mUidTagComplete;
+
+            private NetworkStatsCollection getUidComplete() {
+                if (mUidComplete == null) {
+                    mUidComplete = mUidRecorder.getOrLoadCompleteLocked();
+                }
+                return mUidComplete;
+            }
+
+            private NetworkStatsCollection getUidTagComplete() {
+                if (mUidTagComplete == null) {
+                    mUidTagComplete = mUidTagRecorder.getOrLoadCompleteLocked();
+                }
+                return mUidTagComplete;
+            }
+
+            @Override
+            public NetworkStats getSummaryForNetwork(
+                    NetworkTemplate template, long start, long end) {
+                return mDevStatsCached.getSummary(template, start, end);
+            }
+
+            @Override
+            public NetworkStatsHistory getHistoryForNetwork(NetworkTemplate template, int fields) {
+                return mDevStatsCached.getHistory(template, UID_ALL, SET_ALL, TAG_NONE, fields);
+            }
+
+            @Override
+            public NetworkStats getSummaryForAllUid(
+                    NetworkTemplate template, long start, long end, boolean includeTags) {
+                final NetworkStats stats = getUidComplete().getSummary(template, start, end);
+                if (includeTags) {
+                    final NetworkStats tagStats = getUidTagComplete()
+                            .getSummary(template, start, end);
+                    stats.combineAllValues(tagStats);
+                }
+                return stats;
+            }
+
+            @Override
+            public NetworkStatsHistory getHistoryForUid(
+                    NetworkTemplate template, int uid, int set, int tag, int fields) {
+                if (tag == TAG_NONE) {
+                    return getUidComplete().getHistory(template, uid, set, tag, fields);
+                } else {
+                    return getUidTagComplete().getHistory(template, uid, set, tag, fields);
+                }
+            }
+
+            @Override
+            public void close() {
+                mUidComplete = null;
+                mUidTagComplete = null;
+            }
+        };
     }
 
     @Override
-    public NetworkStats getSummaryForNetwork(NetworkTemplate template, long start, long end) {
-        return mDevStatsCached.getSummary(template, start, end);
-    }
-
-    @Override
-    public NetworkStatsHistory getHistoryForUid(
-            NetworkTemplate template, int uid, int set, int tag, int fields) {
-        // TODO: transition to stats sessions to avoid WeakReferences
-        if (tag == TAG_NONE) {
-            return mUidRecorder.getOrLoadCompleteLocked().getHistory(
-                    template, uid, set, tag, fields);
-        } else {
-            return mUidTagRecorder.getOrLoadCompleteLocked().getHistory(
-                    template, uid, set, tag, fields);
-        }
-    }
-
-    @Override
-    public NetworkStats getSummaryForAllUid(
-            NetworkTemplate template, long start, long end, boolean includeTags) {
-        // TODO: transition to stats sessions to avoid WeakReferences
-        final NetworkStats stats = mUidRecorder.getOrLoadCompleteLocked().getSummary(
-                template, start, end);
-        if (includeTags) {
-            final NetworkStats tagStats = mUidTagRecorder.getOrLoadCompleteLocked().getSummary(
-                    template, start, end);
-            stats.combineAllValues(tagStats);
-        }
-        return stats;
+    public long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
+        mContext.enforceCallingOrSelfPermission(READ_NETWORK_USAGE_HISTORY, TAG);
+        return mDevStatsCached.getSummary(template, start, end).getTotalBytes();
     }
 
     @Override
diff --git a/services/tests/servicestests/src/com/android/server/NetworkPolicyManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/NetworkPolicyManagerServiceTest.java
index 1773e33..2033db6 100644
--- a/services/tests/servicestests/src/com/android/server/NetworkPolicyManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/NetworkPolicyManagerServiceTest.java
@@ -70,7 +70,6 @@
 import android.os.INetworkManagementService;
 import android.os.IPowerManager;
 import android.os.MessageQueue.IdleHandler;
-import android.os.SystemClock;
 import android.os.UserId;
 import android.test.AndroidTestCase;
 import android.test.mock.MockPackageManager;
@@ -628,8 +627,8 @@
         // pretend that 512 bytes total have happened
         stats = new NetworkStats(getElapsedRealtime(), 1)
                 .addIfaceValues(TEST_IFACE, 256L, 2L, 256L, 2L);
-        expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, TIME_MAR_10))
-                .andReturn(stats).atLeastOnce();
+        expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, TIME_MAR_10))
+                .andReturn(stats.getTotalBytes()).atLeastOnce();
         expectPolicyDataEnable(TYPE_WIFI, true);
 
         // TODO: consider making strongly ordered mock
@@ -699,8 +698,8 @@
         {
             expectCurrentTime();
             expect(mConnManager.getAllNetworkState()).andReturn(state).atLeastOnce();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, true);
 
             expectClearNotifications();
@@ -722,8 +721,8 @@
         {
             expectCurrentTime();
             expect(mConnManager.getAllNetworkState()).andReturn(state).atLeastOnce();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, true);
 
             expectRemoveInterfaceQuota(TEST_IFACE);
@@ -745,8 +744,8 @@
 
         {
             expectCurrentTime();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, true);
 
             expectForceUpdate();
@@ -766,8 +765,8 @@
 
         {
             expectCurrentTime();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, false);
 
             expectForceUpdate();
@@ -786,8 +785,8 @@
         {
             expectCurrentTime();
             expect(mConnManager.getAllNetworkState()).andReturn(state).atLeastOnce();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, true);
 
             // snoozed interface still has high quota so background data is
@@ -827,8 +826,8 @@
         {
             expectCurrentTime();
             expect(mConnManager.getAllNetworkState()).andReturn(state).atLeastOnce();
-            expect(mStatsService.getSummaryForNetwork(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
-                    .andReturn(stats).atLeastOnce();
+            expect(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15, currentTimeMillis()))
+                    .andReturn(stats.getTotalBytes()).atLeastOnce();
             expectPolicyDataEnable(TYPE_WIFI, true);
 
             expectRemoveInterfaceQuota(TEST_IFACE);
diff --git a/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java b/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
index 103d8e1..6d9bb29 100644
--- a/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/NetworkStatsServiceTest.java
@@ -54,6 +54,7 @@
 import android.content.Intent;
 import android.net.IConnectivityManager;
 import android.net.INetworkManagementEventObserver;
+import android.net.INetworkStatsSession;
 import android.net.LinkProperties;
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
@@ -84,7 +85,7 @@
  */
 @LargeTest
 public class NetworkStatsServiceTest extends AndroidTestCase {
-    private static final String TAG  = "NetworkStatsServiceTest";
+    private static final String TAG = "NetworkStatsServiceTest";
 
     private static final String TEST_IFACE = "test0";
     private static final String TEST_IFACE2 = "test1";
@@ -113,6 +114,7 @@
     private IConnectivityManager mConnManager;
 
     private NetworkStatsService mService;
+    private INetworkStatsSession mSession;
     private INetworkManagementEventObserver mNetworkObserver;
 
     @Override
@@ -134,6 +136,7 @@
         mService = new NetworkStatsService(
                 mServiceContext, mNetManager, mAlarmManager, mTime, mStatsDir, mSettings);
         mService.bindConnectivityManager(mConnManager);
+        mSession = mService.openSession();
 
         mElapsedRealtime = 0L;
 
@@ -172,6 +175,7 @@
         mSettings = null;
         mConnManager = null;
 
+        mSession.close();
         mService = null;
 
         super.tearDown();
@@ -349,7 +353,7 @@
         mServiceContext.sendBroadcast(new Intent(ACTION_NETWORK_STATS_POLL));
 
         // verify service recorded history
-        history = mService.getHistoryForNetwork(sTemplateWifi, FIELD_ALL);
+        history = mSession.getHistoryForNetwork(sTemplateWifi, FIELD_ALL);
         assertValues(history, Long.MIN_VALUE, Long.MAX_VALUE, 512L, 4L, 512L, 4L, 0);
         assertEquals(HOUR_IN_MILLIS, history.getBucketDuration());
         assertEquals(2, history.size());
@@ -367,7 +371,7 @@
         mServiceContext.sendBroadcast(new Intent(ACTION_NETWORK_STATS_POLL));
 
         // verify identical stats, but spread across 4 buckets now
-        history = mService.getHistoryForNetwork(sTemplateWifi, FIELD_ALL);
+        history = mSession.getHistoryForNetwork(sTemplateWifi, FIELD_ALL);
         assertValues(history, Long.MIN_VALUE, Long.MAX_VALUE, 512L, 4L, 512L, 4L, 0);
         assertEquals(30 * MINUTE_IN_MILLIS, history.getBucketDuration());
         assertEquals(4, history.size());
@@ -652,7 +656,7 @@
         mServiceContext.sendBroadcast(new Intent(ACTION_NETWORK_STATS_POLL));
 
         // first verify entire history present
-        NetworkStats stats = mService.getSummaryForAllUid(
+        NetworkStats stats = mSession.getSummaryForAllUid(
                 sTemplateWifi, Long.MIN_VALUE, Long.MAX_VALUE, true);
         assertEquals(3, stats.size());
         assertValues(stats, IFACE_ALL, UID_RED, SET_DEFAULT, TAG_NONE, 50L, 5L, 50L, 5L, 1);
@@ -661,7 +665,7 @@
 
         // now verify that recent history only contains one uid
         final long currentTime = currentTimeMillis();
-        stats = mService.getSummaryForAllUid(
+        stats = mSession.getSummaryForAllUid(
                 sTemplateWifi, currentTime - HOUR_IN_MILLIS, currentTime, true);
         assertEquals(1, stats.size());
         assertValues(stats, IFACE_ALL, UID_BLUE, SET_DEFAULT, TAG_NONE, 1024L, 8L, 512L, 4L, 0);
@@ -723,7 +727,7 @@
         assertUidTotal(sTemplateWifi, UID_RED, 160L, 4L, 160L, 4L, 2);
 
         // verify entire history present
-        final NetworkStats stats = mService.getSummaryForAllUid(
+        final NetworkStats stats = mSession.getSummaryForAllUid(
                 sTemplateWifi, Long.MIN_VALUE, Long.MAX_VALUE, true);
         assertEquals(4, stats.size());
         assertValues(stats, IFACE_ALL, UID_RED, SET_DEFAULT, TAG_NONE, 128L, 2L, 128L, 2L, 1);
@@ -775,20 +779,20 @@
     }
 
     private void assertNetworkTotal(NetworkTemplate template, long rxBytes, long rxPackets,
-            long txBytes, long txPackets, int operations) {
-        final NetworkStatsHistory history = mService.getHistoryForNetwork(template, FIELD_ALL);
+            long txBytes, long txPackets, int operations) throws Exception {
+        final NetworkStatsHistory history = mSession.getHistoryForNetwork(template, FIELD_ALL);
         assertValues(history, Long.MIN_VALUE, Long.MAX_VALUE, rxBytes, rxPackets, txBytes,
                 txPackets, operations);
     }
 
     private void assertUidTotal(NetworkTemplate template, int uid, long rxBytes, long rxPackets,
-            long txBytes, long txPackets, int operations) {
+            long txBytes, long txPackets, int operations) throws Exception {
         assertUidTotal(template, uid, SET_ALL, rxBytes, rxPackets, txBytes, txPackets, operations);
     }
 
     private void assertUidTotal(NetworkTemplate template, int uid, int set, long rxBytes,
-            long rxPackets, long txBytes, long txPackets, int operations) {
-        final NetworkStatsHistory history = mService.getHistoryForUid(
+            long rxPackets, long txBytes, long txPackets, int operations) throws Exception {
+        final NetworkStatsHistory history = mSession.getHistoryForUid(
                 template, uid, set, TAG_NONE, FIELD_ALL);
         assertValues(history, Long.MIN_VALUE, Long.MAX_VALUE, rxBytes, rxPackets, txBytes,
                 txPackets, operations);
diff --git a/tests/DataIdleTest/src/com/android/tests/dataidle/DataIdleTest.java b/tests/DataIdleTest/src/com/android/tests/dataidle/DataIdleTest.java
index b7e80d4..919e2b3 100644
--- a/tests/DataIdleTest/src/com/android/tests/dataidle/DataIdleTest.java
+++ b/tests/DataIdleTest/src/com/android/tests/dataidle/DataIdleTest.java
@@ -17,15 +17,16 @@
 
 import android.content.Context;
 import android.net.INetworkStatsService;
+import android.net.INetworkStatsSession;
+import android.net.NetworkStats;
 import android.net.NetworkStats.Entry;
 import android.net.NetworkTemplate;
-import android.net.NetworkStats;
+import android.net.TrafficStats;
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.telephony.TelephonyManager;
 import android.test.InstrumentationTestCase;
-import android.test.InstrumentationTestRunner;
 import android.util.Log;
 
 /**
@@ -71,13 +72,17 @@
      * @param template {link {@link NetworkTemplate} to match.
      */
     private void fetchStats(NetworkTemplate template) {
+        INetworkStatsSession session = null;
         try {
             mStatsService.forceUpdate();
-            NetworkStats stats = mStatsService.getSummaryForAllUid(template, Long.MIN_VALUE,
-                    Long.MAX_VALUE, false);
+            session = mStatsService.openSession();
+            final NetworkStats stats = session.getSummaryForAllUid(
+                    template, Long.MIN_VALUE, Long.MAX_VALUE, false);
             reportStats(stats);
         } catch (RemoteException e) {
             Log.w(LOG_TAG, "Failed to fetch network stats.");
+        } finally {
+            TrafficStats.closeQuietly(session);
         }
     }