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
020package org.apache.james.mime4j.storage;
021
022import java.io.ByteArrayInputStream;
023import java.io.IOException;
024import java.io.InputStream;
025import java.io.SequenceInputStream;
026
027import 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 */
042public 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}