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