1 /* 2 * Copyright (C) 2020 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.internal.util; 18 19 import android.annotation.NonNull; 20 import android.util.CharsetUtils; 21 22 import dalvik.system.VMRuntime; 23 24 import java.io.BufferedOutputStream; 25 import java.io.Closeable; 26 import java.io.DataOutput; 27 import java.io.DataOutputStream; 28 import java.io.Flushable; 29 import java.io.IOException; 30 import java.io.OutputStream; 31 import java.util.HashMap; 32 import java.util.Objects; 33 import java.util.concurrent.atomic.AtomicReference; 34 35 /** 36 * Optimized implementation of {@link DataOutput} which buffers data in memory 37 * before flushing to the underlying {@link OutputStream}. 38 * <p> 39 * Benchmarks have demonstrated this class is 2x more efficient than using a 40 * {@link DataOutputStream} with a {@link BufferedOutputStream}. 41 */ 42 public class FastDataOutput implements DataOutput, Flushable, Closeable { 43 private static final int MAX_UNSIGNED_SHORT = 65_535; 44 45 private static final int BUFFER_SIZE = 32_768; 46 47 private static AtomicReference<FastDataOutput> sOutCache = new AtomicReference<>(); 48 49 private final VMRuntime mRuntime; 50 51 private final byte[] mBuffer; 52 private final long mBufferPtr; 53 private final int mBufferCap; 54 55 private OutputStream mOut; 56 private int mBufferPos; 57 58 /** 59 * Values that have been "interned" by {@link #writeInternedUTF(String)}. 60 */ 61 private final HashMap<String, Short> mStringRefs = new HashMap<>(); 62 FastDataOutput(@onNull OutputStream out, int bufferSize)63 public FastDataOutput(@NonNull OutputStream out, int bufferSize) { 64 mRuntime = VMRuntime.getRuntime(); 65 if (bufferSize < 8) { 66 throw new IllegalArgumentException(); 67 } 68 69 mBuffer = (byte[]) mRuntime.newNonMovableArray(byte.class, bufferSize); 70 mBufferPtr = mRuntime.addressOf(mBuffer); 71 mBufferCap = mBuffer.length; 72 73 setOutput(out); 74 } 75 76 /** 77 * Create a new FastDataOutput object or retrieve one from cache. 78 */ obtain(@onNull OutputStream out)79 public static FastDataOutput obtain(@NonNull OutputStream out) { 80 FastDataOutput instance = sOutCache.getAndSet(null); 81 if (instance != null) { 82 instance.setOutput(out); 83 return instance; 84 } 85 return new FastDataOutput(out, BUFFER_SIZE); 86 } 87 88 /** 89 * Put a FastDataOutput object back into the cache. 90 * You must not touch the object after this call. 91 */ release()92 public void release() { 93 if (mBufferPos > 0) { 94 throw new IllegalStateException("Lingering data, call flush() before releasing."); 95 } 96 97 mOut = null; 98 mBufferPos = 0; 99 mStringRefs.clear(); 100 101 if (mBufferCap == BUFFER_SIZE) { 102 // Try to return to the cache. 103 sOutCache.compareAndSet(null, this); 104 } 105 } 106 107 /** 108 * Re-initializes the object for the new output. 109 */ setOutput(@onNull OutputStream out)110 private void setOutput(@NonNull OutputStream out) { 111 mOut = Objects.requireNonNull(out); 112 mBufferPos = 0; 113 mStringRefs.clear(); 114 } 115 drain()116 private void drain() throws IOException { 117 if (mBufferPos > 0) { 118 mOut.write(mBuffer, 0, mBufferPos); 119 mBufferPos = 0; 120 } 121 } 122 123 @Override flush()124 public void flush() throws IOException { 125 drain(); 126 mOut.flush(); 127 } 128 129 @Override close()130 public void close() throws IOException { 131 mOut.close(); 132 release(); 133 } 134 135 @Override write(int b)136 public void write(int b) throws IOException { 137 writeByte(b); 138 } 139 140 @Override write(byte[] b)141 public void write(byte[] b) throws IOException { 142 write(b, 0, b.length); 143 } 144 145 @Override write(byte[] b, int off, int len)146 public void write(byte[] b, int off, int len) throws IOException { 147 if (mBufferCap < len) { 148 drain(); 149 mOut.write(b, off, len); 150 } else { 151 if (mBufferCap - mBufferPos < len) drain(); 152 System.arraycopy(b, off, mBuffer, mBufferPos, len); 153 mBufferPos += len; 154 } 155 } 156 157 @Override writeUTF(String s)158 public void writeUTF(String s) throws IOException { 159 // Attempt to write directly to buffer space if there's enough room, 160 // otherwise fall back to chunking into place 161 if (mBufferCap - mBufferPos < 2 + s.length()) drain(); 162 163 // Magnitude of this returned value indicates the number of bytes 164 // required to encode the string; sign indicates success/failure 165 int len = CharsetUtils.toModifiedUtf8Bytes(s, mBufferPtr, mBufferPos + 2, mBufferCap); 166 if (Math.abs(len) > MAX_UNSIGNED_SHORT) { 167 throw new IOException("Modified UTF-8 length too large: " + len); 168 } 169 170 if (len >= 0) { 171 // Positive value indicates the string was encoded into the buffer 172 // successfully, so we only need to prefix with length 173 writeShort(len); 174 mBufferPos += len; 175 } else { 176 // Negative value indicates buffer was too small and we need to 177 // allocate a temporary buffer for encoding 178 len = -len; 179 final byte[] tmp = (byte[]) mRuntime.newNonMovableArray(byte.class, len + 1); 180 CharsetUtils.toModifiedUtf8Bytes(s, mRuntime.addressOf(tmp), 0, tmp.length); 181 writeShort(len); 182 write(tmp, 0, len); 183 } 184 } 185 186 /** 187 * Write a {@link String} value with the additional signal that the given 188 * value is a candidate for being canonicalized, similar to 189 * {@link String#intern()}. 190 * <p> 191 * Canonicalization is implemented by writing each unique string value once 192 * the first time it appears, and then writing a lightweight {@code short} 193 * reference when that string is written again in the future. 194 * 195 * @see FastDataInput#readInternedUTF() 196 */ writeInternedUTF(@onNull String s)197 public void writeInternedUTF(@NonNull String s) throws IOException { 198 Short ref = mStringRefs.get(s); 199 if (ref != null) { 200 writeShort(ref); 201 } else { 202 writeShort(MAX_UNSIGNED_SHORT); 203 writeUTF(s); 204 205 // We can only safely intern when we have remaining values; if we're 206 // full we at least sent the string value above 207 ref = (short) mStringRefs.size(); 208 if (ref < MAX_UNSIGNED_SHORT) { 209 mStringRefs.put(s, ref); 210 } 211 } 212 } 213 214 @Override writeBoolean(boolean v)215 public void writeBoolean(boolean v) throws IOException { 216 writeByte(v ? 1 : 0); 217 } 218 219 @Override writeByte(int v)220 public void writeByte(int v) throws IOException { 221 if (mBufferCap - mBufferPos < 1) drain(); 222 mBuffer[mBufferPos++] = (byte) ((v >> 0) & 0xff); 223 } 224 225 @Override writeShort(int v)226 public void writeShort(int v) throws IOException { 227 if (mBufferCap - mBufferPos < 2) drain(); 228 mBuffer[mBufferPos++] = (byte) ((v >> 8) & 0xff); 229 mBuffer[mBufferPos++] = (byte) ((v >> 0) & 0xff); 230 } 231 232 @Override writeChar(int v)233 public void writeChar(int v) throws IOException { 234 writeShort((short) v); 235 } 236 237 @Override writeInt(int v)238 public void writeInt(int v) throws IOException { 239 if (mBufferCap - mBufferPos < 4) drain(); 240 mBuffer[mBufferPos++] = (byte) ((v >> 24) & 0xff); 241 mBuffer[mBufferPos++] = (byte) ((v >> 16) & 0xff); 242 mBuffer[mBufferPos++] = (byte) ((v >> 8) & 0xff); 243 mBuffer[mBufferPos++] = (byte) ((v >> 0) & 0xff); 244 } 245 246 @Override writeLong(long v)247 public void writeLong(long v) throws IOException { 248 if (mBufferCap - mBufferPos < 8) drain(); 249 int i = (int) (v >> 32); 250 mBuffer[mBufferPos++] = (byte) ((i >> 24) & 0xff); 251 mBuffer[mBufferPos++] = (byte) ((i >> 16) & 0xff); 252 mBuffer[mBufferPos++] = (byte) ((i >> 8) & 0xff); 253 mBuffer[mBufferPos++] = (byte) ((i >> 0) & 0xff); 254 i = (int) v; 255 mBuffer[mBufferPos++] = (byte) ((i >> 24) & 0xff); 256 mBuffer[mBufferPos++] = (byte) ((i >> 16) & 0xff); 257 mBuffer[mBufferPos++] = (byte) ((i >> 8) & 0xff); 258 mBuffer[mBufferPos++] = (byte) ((i >> 0) & 0xff); 259 } 260 261 @Override writeFloat(float v)262 public void writeFloat(float v) throws IOException { 263 writeInt(Float.floatToIntBits(v)); 264 } 265 266 @Override writeDouble(double v)267 public void writeDouble(double v) throws IOException { 268 writeLong(Double.doubleToLongBits(v)); 269 } 270 271 @Override writeBytes(String s)272 public void writeBytes(String s) throws IOException { 273 // Callers should use writeUTF() 274 throw new UnsupportedOperationException(); 275 } 276 277 @Override writeChars(String s)278 public void writeChars(String s) throws IOException { 279 // Callers should use writeUTF() 280 throw new UnsupportedOperationException(); 281 } 282 } 283