001/*
002 * Copyright 2002-2020 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package org.springframework.core.io.buffer;
017
018import java.util.ArrayList;
019import java.util.Collection;
020import java.util.List;
021import java.util.function.Predicate;
022
023import reactor.core.publisher.Flux;
024
025/**
026 * Custom {@link List} to collect data buffers with and enforce a
027 * limit on the total number of bytes buffered. For use with "collect" or
028 * other buffering operators in declarative APIs, e.g. {@link Flux}.
029 *
030 * <p>Adding elements increases the byte count and if the limit is exceeded,
031 * {@link DataBufferLimitException} is raised.  {@link #clear()} resets the
032 * count. Remove and set are not supported.
033 *
034 * <p><strong>Note:</strong> This class does not automatically release the
035 * buffers it contains. It is usually preferable to use hooks such as
036 * {@link Flux#doOnDiscard} that also take care of cancel and error signals,
037 * or otherwise {@link #releaseAndClear()} can be used.
038 *
039 * @author Rossen Stoyanchev
040 * @since 5.1.11
041 */
042@SuppressWarnings("serial")
043public class LimitedDataBufferList extends ArrayList<DataBuffer> {
044
045        private final int maxByteCount;
046
047        private int byteCount;
048
049
050        public LimitedDataBufferList(int maxByteCount) {
051                this.maxByteCount = maxByteCount;
052        }
053
054
055        @Override
056        public boolean add(DataBuffer buffer) {
057                updateCount(buffer.readableByteCount());
058                return super.add(buffer);
059        }
060
061        @Override
062        public void add(int index, DataBuffer buffer) {
063                super.add(index, buffer);
064                updateCount(buffer.readableByteCount());
065        }
066
067        @Override
068        public boolean addAll(Collection<? extends DataBuffer> collection) {
069                boolean result = super.addAll(collection);
070                collection.forEach(buffer -> updateCount(buffer.readableByteCount()));
071                return result;
072        }
073
074        @Override
075        public boolean addAll(int index, Collection<? extends DataBuffer> collection) {
076                boolean result = super.addAll(index, collection);
077                collection.forEach(buffer -> updateCount(buffer.readableByteCount()));
078                return result;
079        }
080
081        private void updateCount(int bytesToAdd) {
082                if (this.maxByteCount < 0) {
083                        return;
084                }
085                if (bytesToAdd > Integer.MAX_VALUE - this.byteCount) {
086                        raiseLimitException();
087                }
088                else {
089                        this.byteCount += bytesToAdd;
090                        if (this.byteCount > this.maxByteCount) {
091                                raiseLimitException();
092                        }
093                }
094        }
095
096        private void raiseLimitException() {
097                // Do not release here, it's likely down via doOnDiscard..
098                throw new DataBufferLimitException(
099                                "Exceeded limit on max bytes to buffer : " + this.maxByteCount);
100        }
101
102        @Override
103        public DataBuffer remove(int index) {
104                throw new UnsupportedOperationException();
105        }
106
107        @Override
108        public boolean remove(Object o) {
109                throw new UnsupportedOperationException();
110        }
111
112        @Override
113        protected void removeRange(int fromIndex, int toIndex) {
114                throw new UnsupportedOperationException();
115        }
116
117        @Override
118        public boolean removeAll(Collection<?> c) {
119                throw new UnsupportedOperationException();
120        }
121
122        @Override
123        public boolean removeIf(Predicate<? super DataBuffer> filter) {
124                throw new UnsupportedOperationException();
125        }
126
127        @Override
128        public DataBuffer set(int index, DataBuffer element) {
129                throw new UnsupportedOperationException();
130        }
131
132        @Override
133        public void clear() {
134                this.byteCount = 0;
135                super.clear();
136        }
137
138        /**
139         * Shortcut to {@link DataBufferUtils#release release} all data buffers and
140         * then {@link #clear()}.
141         */
142        public void releaseAndClear() {
143                forEach(buf -> {
144                        try {
145                                DataBufferUtils.release(buf);
146                        }
147                        catch (Throwable ex) {
148                                // Keep going..
149                        }
150                });
151                clear();
152        }
153
154}