diff --git a/core/java/android/net/NetworkStatsHistory.java b/core/java/android/net/NetworkStatsHistory.java
index bf25602..f413063 100644
--- a/core/java/android/net/NetworkStatsHistory.java
+++ b/core/java/android/net/NetworkStatsHistory.java
@@ -45,8 +45,8 @@
 import libcore.util.EmptyArray;
 
 import java.io.CharArrayWriter;
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
+import java.io.DataInput;
+import java.io.DataOutput;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.net.ProtocolException;
@@ -162,7 +162,7 @@
         out.writeLong(totalBytes);
     }
 
-    public NetworkStatsHistory(DataInputStream in) throws IOException {
+    public NetworkStatsHistory(DataInput in) throws IOException {
         final int version = in.readInt();
         switch (version) {
             case VERSION_INIT: {
@@ -204,7 +204,7 @@
         }
     }
 
-    public void writeToStream(DataOutputStream out) throws IOException {
+    public void writeToStream(DataOutput out) throws IOException {
         out.writeInt(VERSION_ADD_ACTIVE);
         out.writeLong(bucketDuration);
         writeVarLongArray(out, bucketStart, bucketCount);
@@ -768,7 +768,7 @@
      */
     public static class DataStreamUtils {
         @Deprecated
-        public static long[] readFullLongArray(DataInputStream in) throws IOException {
+        public static long[] readFullLongArray(DataInput in) throws IOException {
             final int size = in.readInt();
             if (size < 0) throw new ProtocolException("negative array size");
             final long[] values = new long[size];
@@ -781,7 +781,7 @@
         /**
          * Read variable-length {@link Long} using protobuf-style approach.
          */
-        public static long readVarLong(DataInputStream in) throws IOException {
+        public static long readVarLong(DataInput in) throws IOException {
             int shift = 0;
             long result = 0;
             while (shift < 64) {
@@ -797,7 +797,7 @@
         /**
          * Write variable-length {@link Long} using protobuf-style approach.
          */
-        public static void writeVarLong(DataOutputStream out, long value) throws IOException {
+        public static void writeVarLong(DataOutput out, long value) throws IOException {
             while (true) {
                 if ((value & ~0x7FL) == 0) {
                     out.writeByte((int) value);
@@ -809,7 +809,7 @@
             }
         }
 
-        public static long[] readVarLongArray(DataInputStream in) throws IOException {
+        public static long[] readVarLongArray(DataInput in) throws IOException {
             final int size = in.readInt();
             if (size == -1) return null;
             if (size < 0) throw new ProtocolException("negative array size");
@@ -820,7 +820,7 @@
             return values;
         }
 
-        public static void writeVarLongArray(DataOutputStream out, long[] values, int size)
+        public static void writeVarLongArray(DataOutput out, long[] values, int size)
                 throws IOException {
             if (values == null) {
                 out.writeInt(-1);
diff --git a/services/core/java/com/android/server/net/NetworkIdentitySet.java b/services/core/java/com/android/server/net/NetworkIdentitySet.java
index 2326ad3..bce8069 100644
--- a/services/core/java/com/android/server/net/NetworkIdentitySet.java
+++ b/services/core/java/com/android/server/net/NetworkIdentitySet.java
@@ -20,8 +20,8 @@
 import android.service.NetworkIdentitySetProto;
 import android.util.proto.ProtoOutputStream;
 
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
+import java.io.DataInput;
+import java.io.DataOutput;
 import java.io.IOException;
 import java.util.HashSet;
 
@@ -44,7 +44,7 @@
     public NetworkIdentitySet() {
     }
 
-    public NetworkIdentitySet(DataInputStream in) throws IOException {
+    public NetworkIdentitySet(DataInput in) throws IOException {
         final int version = in.readInt();
         final int size = in.readInt();
         for (int i = 0; i < size; i++) {
@@ -89,7 +89,7 @@
         }
     }
 
-    public void writeToStream(DataOutputStream out) throws IOException {
+    public void writeToStream(DataOutput out) throws IOException {
         out.writeInt(VERSION_ADD_DEFAULT_NETWORK);
         out.writeInt(size());
         for (NetworkIdentity ident : this) {
@@ -143,7 +143,7 @@
         return true;
     }
 
-    private static void writeOptionalString(DataOutputStream out, String value) throws IOException {
+    private static void writeOptionalString(DataOutput out, String value) throws IOException {
         if (value != null) {
             out.writeByte(1);
             out.writeUTF(value);
@@ -152,7 +152,7 @@
         }
     }
 
-    private static String readOptionalString(DataInputStream in) throws IOException {
+    private static String readOptionalString(DataInput in) throws IOException {
         if (in.readByte() != 0) {
             return in.readUTF();
         } else {
diff --git a/services/core/java/com/android/server/net/NetworkStatsCollection.java b/services/core/java/com/android/server/net/NetworkStatsCollection.java
index c4beddd..6aefe41 100644
--- a/services/core/java/com/android/server/net/NetworkStatsCollection.java
+++ b/services/core/java/com/android/server/net/NetworkStatsCollection.java
@@ -63,12 +63,15 @@
 import com.google.android.collect.Maps;
 
 import java.io.BufferedInputStream;
+import java.io.DataInput;
 import java.io.DataInputStream;
+import java.io.DataOutput;
 import java.io.DataOutputStream;
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.PrintWriter;
 import java.net.ProtocolException;
 import java.time.ZonedDateTime;
@@ -82,7 +85,7 @@
  * Collection of {@link NetworkStatsHistory}, stored based on combined key of
  * {@link NetworkIdentitySet}, UID, set, and tag. Knows how to persist itself.
  */
-public class NetworkStatsCollection implements FileRotator.Reader {
+public class NetworkStatsCollection implements FileRotator.Reader, FileRotator.Writer {
     /** File header magic number: "ANET" */
     private static final int FILE_MAGIC = 0x414E4554;
 
@@ -431,10 +434,10 @@
 
     @Override
     public void read(InputStream in) throws IOException {
-        read(new DataInputStream(in));
+        read((DataInput) new DataInputStream(in));
     }
 
-    public void read(DataInputStream in) throws IOException {
+    private void read(DataInput in) throws IOException {
         // verify file magic header intact
         final int magic = in.readInt();
         if (magic != FILE_MAGIC) {
@@ -468,7 +471,13 @@
         }
     }
 
-    public void write(DataOutputStream out) throws IOException {
+    @Override
+    public void write(OutputStream out) throws IOException {
+        write((DataOutput) new DataOutputStream(out));
+        out.flush();
+    }
+
+    private void write(DataOutput out) throws IOException {
         // cluster key lists grouped by ident
         final HashMap<NetworkIdentitySet, ArrayList<Key>> keysByIdent = Maps.newHashMap();
         for (Key key : mStats.keySet()) {
@@ -497,8 +506,6 @@
                 history.writeToStream(out);
             }
         }
-
-        out.flush();
     }
 
     @Deprecated
diff --git a/services/core/java/com/android/server/net/NetworkStatsRecorder.java b/services/core/java/com/android/server/net/NetworkStatsRecorder.java
index ce74169..978ae87 100644
--- a/services/core/java/com/android/server/net/NetworkStatsRecorder.java
+++ b/services/core/java/com/android/server/net/NetworkStatsRecorder.java
@@ -42,7 +42,6 @@
 import libcore.io.IoUtils;
 
 import java.io.ByteArrayOutputStream;
-import java.io.DataOutputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -375,7 +374,7 @@
 
         @Override
         public void write(OutputStream out) throws IOException {
-            mCollection.write(new DataOutputStream(out));
+            mCollection.write(out);
             mCollection.reset();
         }
     }
@@ -412,7 +411,7 @@
 
         @Override
         public void write(OutputStream out) throws IOException {
-            mTemp.write(new DataOutputStream(out));
+            mTemp.write(out);
         }
     }