001    /**
002     * Copyright (c) 2000-2013 Liferay, Inc. All rights reserved.
003     *
004     * This library is free software; you can redistribute it and/or modify it under
005     * the terms of the GNU Lesser General Public License as published by the Free
006     * Software Foundation; either version 2.1 of the License, or (at your option)
007     * any later version.
008     *
009     * This library is distributed in the hope that it will be useful, but WITHOUT
010     * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
011     * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
012     * details.
013     */
014    
015    package com.liferay.portal.kernel.io;
016    
017    import java.io.IOException;
018    import java.io.OutputStream;
019    
020    import java.nio.ByteBuffer;
021    
022    /**
023     * @author Shuyang Zhou
024     */
025    public class RestrictedByteArrayCacheOutputStream extends OutputStream {
026    
027            public RestrictedByteArrayCacheOutputStream(
028                    OutputStream outputStream, int cacheCapacity,
029                    FlushPreAction flushPreAction) {
030    
031                    this(outputStream, 32, cacheCapacity, flushPreAction);
032            }
033    
034            public RestrictedByteArrayCacheOutputStream(
035                    OutputStream outputStream, int initialCacheSize, int cacheCapacity,
036                    FlushPreAction flushPreAction) {
037    
038                    if (initialCacheSize > cacheCapacity) {
039                            throw new IllegalArgumentException(
040                                    "Initial cache size " + initialCacheSize +
041                                            " is larger than cache capacity " + cacheCapacity);
042                    }
043    
044                    this.outputStream = outputStream;
045                    this.cacheCapacity = cacheCapacity;
046                    this.flushPreAction = flushPreAction;
047    
048                    cache = new byte[initialCacheSize];
049            }
050    
051            @Override
052            public void flush() throws IOException {
053                    if (overflowed) {
054                            return;
055                    }
056    
057                    if (flushPreAction != null) {
058                            flushPreAction.beforeFlush();
059                    }
060    
061                    overflowed = true;
062    
063                    outputStream.write(cache, 0, index);
064    
065                    cache = null;
066                    index = -1;
067            }
068    
069            public int getCacheCapacity() {
070                    return cacheCapacity;
071            }
072    
073            public boolean isOverflowed() {
074                    return overflowed;
075            }
076    
077            public void reset() {
078                    if (overflowed) {
079                            throw new IllegalStateException("Cache overflowed");
080                    }
081    
082                    index = 0;
083            }
084    
085            public int size() {
086                    return index;
087            }
088    
089            public byte[] toByteArray() {
090                    if (overflowed) {
091                            throw new IllegalStateException("Cache overflowed");
092                    }
093    
094                    byte[] newCache = new byte[index];
095    
096                    System.arraycopy(cache, 0, newCache, 0, index);
097    
098                    return newCache;
099            }
100    
101            public byte[] unsafeGetByteArray() {
102                    if (overflowed) {
103                            throw new IllegalStateException("Cache overflowed");
104                    }
105    
106                    return cache;
107            }
108    
109            public ByteBuffer unsafeGetByteBuffer() {
110                    if (overflowed) {
111                            throw new IllegalStateException("Cache overflowed");
112                    }
113    
114                    return ByteBuffer.wrap(cache, 0, index);
115            }
116    
117            @Override
118            public void write(byte[] bytes) throws IOException {
119                    write(bytes, 0, bytes.length);
120            }
121    
122            @Override
123            public void write(byte[] bytes, int offset, int length) throws IOException {
124                    if (length <= 0) {
125                            return;
126                    }
127    
128                    if (overflowed) {
129                            outputStream.write(bytes, offset, length);
130    
131                            return;
132                    }
133    
134                    int newIndex = index + length;
135    
136                    if (newIndex > cacheCapacity) {
137                            flush();
138    
139                            outputStream.write(bytes, offset, length);
140    
141                            return;
142                    }
143    
144                    ensureCacheSize(newIndex);
145    
146                    System.arraycopy(bytes, offset, cache, index, length);
147    
148                    index = newIndex;
149            }
150    
151            @Override
152            public void write(int b) throws IOException {
153                    if (overflowed) {
154                            outputStream.write(b);
155    
156                            return;
157                    }
158    
159                    int newIndex = index + 1;
160    
161                    if (newIndex > cacheCapacity) {
162                            flush();
163    
164                            outputStream.write(b);
165    
166                            return;
167                    }
168    
169                    ensureCacheSize(newIndex);
170    
171                    cache[index] = (byte)b;
172    
173                    index = newIndex;
174            }
175    
176            public static interface FlushPreAction {
177    
178                    public void beforeFlush() throws IOException;
179    
180            }
181    
182            protected void ensureCacheSize(int newIndex) {
183                    if (newIndex <= cache.length) {
184                            return;
185                    }
186    
187                    int newCacheSize = Math.max(cache.length << 1, newIndex);
188    
189                    if (newCacheSize > cacheCapacity) {
190                            newCacheSize = cacheCapacity;
191                    }
192    
193                    byte[] newCache = new byte[newCacheSize];
194    
195                    System.arraycopy(cache, 0, newCache, 0, cache.length);
196    
197                    cache = newCache;
198            }
199    
200            protected byte[] cache;
201            protected int cacheCapacity;
202            protected FlushPreAction flushPreAction;
203            protected int index;
204            protected OutputStream outputStream;
205            protected boolean overflowed;
206    
207    }