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