001/*
002 * Copyright 2024-2025 Revetware LLC.
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 * http://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 */
016
017package com.soklet.servlet.javax;
018
019import javax.annotation.Nonnull;
020import javax.annotation.Nullable;
021import javax.annotation.concurrent.NotThreadSafe;
022import javax.servlet.ServletOutputStream;
023import javax.servlet.WriteListener;
024import java.io.IOException;
025import java.io.OutputStream;
026import java.util.function.BiConsumer;
027import java.util.function.Consumer;
028
029import static java.lang.String.format;
030import static java.util.Objects.requireNonNull;
031
032/**
033 * Soklet integration implementation of {@link ServletOutputStream}.
034 *
035 * @author <a href="https://www.revetkn.com">Mark Allen</a>
036 */
037@NotThreadSafe
038public final class SokletServletOutputStream extends ServletOutputStream {
039        @Nonnull
040        private final OutputStream outputStream;
041        @Nonnull
042        private final BiConsumer<SokletServletOutputStream, Integer> onWriteOccurred;
043        @Nonnull
044        private final Consumer<SokletServletOutputStream> onWriteFinalized;
045        @Nonnull
046        private Boolean writeFinalized;
047
048        @Nonnull
049        public static Builder withOutputStream(@Nonnull OutputStream outputStream) {
050                return new Builder(outputStream);
051        }
052
053        private SokletServletOutputStream(@Nonnull Builder builder) {
054                requireNonNull(builder);
055                requireNonNull(builder.outputStream);
056
057                this.outputStream = builder.outputStream;
058                this.onWriteOccurred = builder.onWriteOccurred != null ? builder.onWriteOccurred : (ignored1, ignored2) -> {};
059                this.onWriteFinalized = builder.onWriteFinalized != null ? builder.onWriteFinalized : (ignored) -> {};
060                this.writeFinalized = false;
061        }
062
063        /**
064         * Builder used to construct instances of {@link SokletServletOutputStream}.
065         * <p>
066         * This class is intended for use by a single thread.
067         *
068         * @author <a href="https://www.revetkn.com">Mark Allen</a>
069         */
070        @NotThreadSafe
071        public static class Builder {
072                @Nonnull
073                private OutputStream outputStream;
074                @Nullable
075                private BiConsumer<SokletServletOutputStream, Integer> onWriteOccurred;
076                @Nullable
077                private Consumer<SokletServletOutputStream> onWriteFinalized;
078
079                @Nonnull
080                private Builder(@Nonnull OutputStream outputStream) {
081                        requireNonNull(outputStream);
082                        this.outputStream = outputStream;
083                }
084
085                @Nonnull
086                public Builder outputStream(@Nonnull OutputStream outputStream) {
087                        requireNonNull(outputStream);
088                        this.outputStream = outputStream;
089                        return this;
090                }
091
092                @Nonnull
093                public Builder onWriteOccurred(@Nullable BiConsumer<SokletServletOutputStream, Integer> onWriteOccurred) {
094                        this.onWriteOccurred = onWriteOccurred;
095                        return this;
096                }
097
098                @Nonnull
099                public Builder onWriteFinalized(@Nullable Consumer<SokletServletOutputStream> onWriteFinalized) {
100                        this.onWriteFinalized = onWriteFinalized;
101                        return this;
102                }
103
104                @Nonnull
105                public SokletServletOutputStream build() {
106                        return new SokletServletOutputStream(this);
107                }
108        }
109
110        @Nonnull
111        protected OutputStream getOutputStream() {
112                return this.outputStream;
113        }
114
115        @Nonnull
116        protected BiConsumer<SokletServletOutputStream, Integer> getOnWriteOccurred() {
117                return this.onWriteOccurred;
118        }
119
120        @Nonnull
121        protected Consumer<SokletServletOutputStream> getOnWriteFinalized() {
122                return this.onWriteFinalized;
123        }
124
125        @Nonnull
126        protected Boolean getWriteFinalized() {
127                return this.writeFinalized;
128        }
129
130        protected void setWriteFinalized(@Nonnull Boolean writeFinalized) {
131                requireNonNull(writeFinalized);
132                this.writeFinalized = writeFinalized;
133        }
134
135// Implementation of ServletOutputStream methods below:
136
137        @Override
138        public void write(int b) throws IOException {
139                getOutputStream().write(b);
140                getOnWriteOccurred().accept(this, b);
141        }
142
143        @Override
144        public boolean isReady() {
145                return !getWriteFinalized();
146        }
147
148        @Override
149        public void flush() throws IOException {
150                super.flush();
151                getOutputStream().flush();
152
153                if (!getWriteFinalized()) {
154                        setWriteFinalized(true);
155                        getOnWriteFinalized().accept(this);
156                }
157        }
158
159        @Override
160        public void close() throws IOException {
161                super.close();
162                getOutputStream().close();
163
164                if (!getWriteFinalized()) {
165                        setWriteFinalized(true);
166                        getOnWriteFinalized().accept(this);
167                }
168        }
169
170        @Override
171        public void setWriteListener(@Nonnull WriteListener writeListener) {
172                requireNonNull(writeListener);
173                throw new IllegalStateException(format("%s functionality is not supported", WriteListener.class.getSimpleName()));
174        }
175}