1 /* 2 * Copyright (C) 2013 The Android Open Source Project 3 * 4 * Licensed under the Apache 5 * License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 */ 17 18 package com.android.tradefed.util; 19 20 import com.google.common.io.CountingOutputStream; 21 22 import java.io.BufferedOutputStream; 23 import java.io.ByteArrayInputStream; 24 import java.io.File; 25 import java.io.FileInputStream; 26 import java.io.FileNotFoundException; 27 import java.io.FileOutputStream; 28 import java.io.IOException; 29 import java.io.InputStream; 30 import java.io.OutputStream; 31 import java.io.SequenceInputStream; 32 33 /** 34 * A thread safe file backed {@link OutputStream} that limits the maximum amount of data that can be 35 * written. 36 * <p/> 37 * This is implemented by keeping a circular list of Files of fixed size. Once a File has reached a 38 * certain size, the class jumps to use the next File in the list. If the next File is non empty, it 39 * is deleted, and a new file created. 40 */ 41 public class SizeLimitedOutputStream extends OutputStream { 42 43 private static final int DEFAULT_NUM_TMP_FILES = 5; 44 45 /** The max number of bytes to store in the buffer */ 46 private static final int BUFF_SIZE = 32 * 1024; 47 48 // circular array of backing files 49 private final File[] mFiles; 50 private final long mMaxFileSize; 51 private CountingOutputStream mCurrentOutputStream; 52 private int mCurrentFilePos = 0; 53 private final String mTempFilePrefix; 54 private final String mTempFileSuffix; 55 56 /** 57 * Creates a {@link SizeLimitedOutputStream}. 58 * 59 * @param maxDataSize the approximate max size in bytes to keep in the output stream 60 * @param numFiles the max number of backing files to use to store data. Higher values will mean 61 * max data kept will be close to maxDataSize, but with a possible performance 62 * penalty. 63 * @param tempFilePrefix prefix to use for temporary files 64 * @param tempFileSuffix suffix to use for temporary files 65 */ 66 public SizeLimitedOutputStream(long maxDataSize, int numFiles, String tempFilePrefix, 67 String tempFileSuffix) { 68 mMaxFileSize = maxDataSize / numFiles; 69 mFiles = new File[numFiles]; 70 mCurrentFilePos = numFiles; 71 mTempFilePrefix = tempFilePrefix; 72 mTempFileSuffix = tempFileSuffix; 73 } 74 75 /** 76 * Creates a {@link SizeLimitedOutputStream} with default number of backing files. 77 * 78 * @param maxDataSize the approximate max size to keep in the output stream 79 * @param tempFilePrefix prefix to use for temporary files 80 * @param tempFileSuffix suffix to use for temporary files 81 */ 82 public SizeLimitedOutputStream(long maxDataSize, String tempFilePrefix, String tempFileSuffix) { 83 this(maxDataSize, DEFAULT_NUM_TMP_FILES, tempFilePrefix, tempFileSuffix); 84 } 85 86 /** 87 * Gets the collected output as a {@link InputStream}. 88 * <p/> 89 * It is recommended to buffer returned stream before using. 90 * 91 * @return The collected output as a {@link InputStream}. 92 */ 93 public synchronized InputStream getData() throws IOException { 94 flush(); 95 InputStream combinedStream = null; 96 for (int i = 0; i < mFiles.length; i++) { 97 // oldest/starting file is always the next one up from current 98 int currentPos = (mCurrentFilePos + i + 1) % mFiles.length; 99 if (mFiles[currentPos] != null) { 100 @SuppressWarnings("resource") 101 FileInputStream fStream = new FileInputStream(mFiles[currentPos]); 102 if (combinedStream == null) { 103 combinedStream = fStream; 104 } else { 105 combinedStream = new SequenceInputStream(combinedStream, fStream); 106 } 107 } 108 } 109 if (combinedStream == null) { 110 combinedStream = new ByteArrayInputStream(new byte[0]); 111 } 112 return combinedStream; 113 114 } 115 116 /** 117 * {@inheritDoc} 118 */ 119 @Override 120 public synchronized void flush() { 121 if (mCurrentOutputStream == null) { 122 return; 123 } 124 try { 125 mCurrentOutputStream.flush(); 126 } catch (IOException e) { 127 // don't use CLog in this class, because its the underlying stream for the logger. 128 // leads to bad things 129 System.out.printf("failed to flush data: %s\n", e); 130 } 131 } 132 133 /** 134 * Delete all accumulated data. 135 */ 136 public void delete() { 137 close(); 138 for (int i = 0; i < mFiles.length; i++) { 139 FileUtil.deleteFile(mFiles[i]); 140 mFiles[i] = null; 141 } 142 } 143 144 /** 145 * Closes the write stream 146 */ 147 @Override 148 public synchronized void close() { 149 StreamUtil.flushAndCloseStream(mCurrentOutputStream); 150 mCurrentOutputStream = null; 151 } 152 153 /** 154 * Creates a new tmp file, closing the old one as necessary 155 * <p> 156 * Exposed for unit testing. 157 * </p> 158 * 159 * @throws IOException 160 * @throws FileNotFoundException 161 */ 162 synchronized void generateNextFile() throws IOException, FileNotFoundException { 163 // close current stream 164 close(); 165 mCurrentFilePos = getNextIndex(mCurrentFilePos); 166 FileUtil.deleteFile(mFiles[mCurrentFilePos]); 167 mFiles[mCurrentFilePos] = FileUtil.createTempFile(mTempFilePrefix, mTempFileSuffix); 168 mCurrentOutputStream = new CountingOutputStream(new BufferedOutputStream( 169 new FileOutputStream(mFiles[mCurrentFilePos]), BUFF_SIZE)); 170 } 171 172 /** 173 * Gets the next index to use for <var>mFiles</var>, treating it as a circular list. 174 */ 175 private int getNextIndex(int i) { 176 return (i + 1) % mFiles.length; 177 } 178 179 @Override 180 public synchronized void write(int data) throws IOException { 181 if (mCurrentOutputStream == null) { 182 generateNextFile(); 183 } 184 mCurrentOutputStream.write(data); 185 if (mCurrentOutputStream.getCount() >= mMaxFileSize) { 186 generateNextFile(); 187 } 188 } 189 190 @Override 191 public synchronized void write(byte[] b, int off, int len) throws IOException { 192 if (mCurrentOutputStream == null) { 193 generateNextFile(); 194 } 195 // keep writing to output stream as long as we have something to write 196 while (len > 0) { 197 // get current output stream size 198 long currentSize = mCurrentOutputStream.getCount(); 199 // get how many more we can write into current 200 long freeSpace = mMaxFileSize - currentSize; 201 // decide how much we should write: either fill up free space, or write entire content 202 long sizeToWrite = freeSpace > len ? len : freeSpace; 203 mCurrentOutputStream.write(b, off, (int)sizeToWrite); 204 // accounting of space left, where to write next 205 freeSpace -= sizeToWrite; 206 off += sizeToWrite; 207 len -= sizeToWrite; 208 if (freeSpace <= 0) { 209 generateNextFile(); 210 } 211 } 212 } 213 } 214