001    /****************************************************************
002     * Licensed to the Apache Software Foundation (ASF) under one   *
003     * or more contributor license agreements.  See the NOTICE file *
004     * distributed with this work for additional information        *
005     * regarding copyright ownership.  The ASF licenses this file   *
006     * to you under the Apache License, Version 2.0 (the            *
007     * "License"); you may not use this file except in compliance   *
008     * with the License.  You may obtain a copy of the License at   *
009     *                                                              *
010     *   http://www.apache.org/licenses/LICENSE-2.0                 *
011     *                                                              *
012     * Unless required by applicable law or agreed to in writing,   *
013     * software distributed under the License is distributed on an  *
014     * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY       *
015     * KIND, either express or implied.  See the License for the    *
016     * specific language governing permissions and limitations      *
017     * under the License.                                           *
018     ****************************************************************/
019    
020    package org.apache.james.mime4j.storage;
021    
022    import java.io.ByteArrayInputStream;
023    import java.io.IOException;
024    import java.io.InputStream;
025    import java.io.SequenceInputStream;
026    
027    import org.apache.james.mime4j.util.ByteArrayBuffer;
028    
029    /**
030     * A {@link StorageProvider} that keeps small amounts of data in memory and
031     * writes the remainder to another <code>StorageProvider</code> (the back-end)
032     * if a certain threshold size gets exceeded.
033     * <p>
034     * Example usage:
035     *
036     * <pre>
037     * StorageProvider tempStore = new TempFileStorageProvider();
038     * StorageProvider provider = new ThresholdStorageProvider(tempStore, 4096);
039     * DefaultStorageProvider.setInstance(provider);
040     * </pre>
041     */
042    public class ThresholdStorageProvider extends AbstractStorageProvider {
043    
044        private final StorageProvider backend;
045        private final int thresholdSize;
046    
047        /**
048         * Creates a new <code>ThresholdStorageProvider</code> for the given
049         * back-end using a threshold size of 2048 bytes.
050         */
051        public ThresholdStorageProvider(StorageProvider backend) {
052            this(backend, 2048);
053        }
054    
055        /**
056         * Creates a new <code>ThresholdStorageProvider</code> for the given
057         * back-end and threshold size.
058         *
059         * @param backend
060         *            used to store the remainder of the data if the threshold size
061         *            gets exceeded.
062         * @param thresholdSize
063         *            determines how much bytes are kept in memory before that
064         *            back-end storage provider is used to store the remainder of
065         *            the data.
066         */
067        public ThresholdStorageProvider(StorageProvider backend, int thresholdSize) {
068            if (backend == null)
069                throw new IllegalArgumentException();
070            if (thresholdSize < 1)
071                throw new IllegalArgumentException();
072    
073            this.backend = backend;
074            this.thresholdSize = thresholdSize;
075        }
076    
077        public StorageOutputStream createStorageOutputStream() {
078            return new ThresholdStorageOutputStream();
079        }
080    
081        private final class ThresholdStorageOutputStream extends
082                StorageOutputStream {
083    
084            private final ByteArrayBuffer head;
085            private StorageOutputStream tail;
086    
087            public ThresholdStorageOutputStream() {
088                final int bufferSize = Math.min(thresholdSize, 1024);
089                head = new ByteArrayBuffer(bufferSize);
090            }
091    
092            @Override
093            public void close() throws IOException {
094                super.close();
095    
096                if (tail != null)
097                    tail.close();
098            }
099    
100            @Override
101            protected void write0(byte[] buffer, int offset, int length)
102                    throws IOException {
103                int remainingHeadSize = thresholdSize - head.length();
104                if (remainingHeadSize > 0) {
105                    int n = Math.min(remainingHeadSize, length);
106                    head.append(buffer, offset, n);
107                    offset += n;
108                    length -= n;
109                }
110    
111                if (length > 0) {
112                    if (tail == null)
113                        tail = backend.createStorageOutputStream();
114    
115                    tail.write(buffer, offset, length);
116                }
117            }
118    
119            @Override
120            protected Storage toStorage0() throws IOException {
121                if (tail == null)
122                    return new MemoryStorageProvider.MemoryStorage(head.buffer(),
123                            head.length());
124    
125                return new ThresholdStorage(head.buffer(), head.length(), tail
126                        .toStorage());
127            }
128    
129        }
130    
131        private static final class ThresholdStorage implements Storage {
132    
133            private byte[] head;
134            private final int headLen;
135            private Storage tail;
136    
137            public ThresholdStorage(byte[] head, int headLen, Storage tail) {
138                this.head = head;
139                this.headLen = headLen;
140                this.tail = tail;
141            }
142    
143            public void delete() {
144                if (head != null) {
145                    head = null;
146                    tail.delete();
147                    tail = null;
148                }
149            }
150    
151            public InputStream getInputStream() throws IOException {
152                if (head == null)
153                    throw new IllegalStateException("storage has been deleted");
154    
155                InputStream headStream = new ByteArrayInputStream(head, 0, headLen);
156                InputStream tailStream = tail.getInputStream();
157                return new SequenceInputStream(headStream, tailStream);
158            }
159    
160        }
161    }