001/*
002 * Copyright (C) 2022 - 2024, 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 *    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 */
016package io.github.ascopes.jct.diagnostics;
017
018import static java.util.Objects.requireNonNull;
019
020import java.io.IOException;
021import java.io.OutputStream;
022import java.io.OutputStreamWriter;
023import java.io.Writer;
024import java.nio.charset.Charset;
025import java.util.concurrent.locks.Lock;
026import java.util.concurrent.locks.ReentrantLock;
027import org.apiguardian.api.API;
028import org.apiguardian.api.API.Status;
029
030/**
031 * A writer that wraps an output stream and also writes any content to an in-memory buffer.
032 *
033 * <p>This is thread-safe.
034 *
035 * @author Ashley Scopes
036 * @since 0.0.1
037 */
038@API(since = "0.0.1", status = Status.STABLE)
039public final class TeeWriter extends Writer {
040
041  private final Lock lock;
042
043  private volatile boolean closed;
044
045  private final Writer writer;
046
047  // We use a StringBuilder and manually synchronise it rather than
048  // a string buffer, as we want to manually synchronise the builder
049  // and the delegated output writer at the same time.
050  private final StringBuilder builder;
051
052  /**
053   * Initialise the writer.
054   *
055   * @param writer the underlying writer to "tee" to.
056   */
057  public TeeWriter(Writer writer) {
058    lock = new ReentrantLock();
059    closed = false;
060
061    this.writer = requireNonNull(writer, "writer");
062    builder = new StringBuilder();
063    builder.ensureCapacity(64);
064  }
065
066  @Override
067  public void close() throws IOException {
068    // release to set and acquire to check ensures in-order operations to prevent
069    // a very minute chance of a race condition.
070    lock.lock();
071    try {
072      if (!closed) {
073        closed = true;
074        builder.trimToSize();
075        writer.flush();
076        writer.close();
077      }
078    } finally {
079      lock.unlock();
080    }
081  }
082
083  @Override
084  public void flush() throws IOException {
085    lock.lock();
086    try {
087      ensureOpen();
088      writer.flush();
089    } finally {
090      lock.unlock();
091    }
092  }
093
094  /**
095   * Get the content of the internal buffer.
096   *
097   * @return the content.
098   * @since 0.2.1
099   */
100  @API(since = "0.2.1", status = Status.STABLE)
101  public String getContent() {
102    lock.lock();
103    try {
104      return builder.toString();
105    } finally {
106      lock.unlock();
107    }
108  }
109
110  /**
111   * Get the content of the internal buffer.
112   *
113   * <p>This calls {@link #getContent()} internally as of 0.2.1.
114   *
115   * @return the content.
116   */
117  @Override
118  public String toString() {
119    return getContent();
120  }
121
122  @Override
123  public void write(char[] buffer, int offset, int length) throws IOException {
124    lock.lock();
125    try {
126      ensureOpen();
127
128      writer.write(buffer, offset, length);
129      // Only append to the buffer once we know that the writing
130      // operation has completed.
131      builder.append(buffer, offset, length);
132    } finally {
133      lock.unlock();
134    }
135  }
136
137  private void ensureOpen() throws IOException {
138    if (closed) {
139      throw new IOException("TeeWriter is closed");
140    }
141  }
142
143  /**
144   * Create a tee writer for the given output stream.
145   *
146   * <p>Remember you may need to manually flush the tee writer for all contents to be committed to
147   * the output stream.
148   *
149   * @param outputStream the output stream.
150   * @param charset      the charset.
151   * @return the Tee Writer.
152   * @since 0.2.1
153   */
154  @API(since = "0.2.1", status = Status.STABLE)
155  public static TeeWriter wrapOutputStream(OutputStream outputStream, Charset charset) {
156    requireNonNull(outputStream, "outputStream");
157    requireNonNull(charset, "charset");
158    var writer = new OutputStreamWriter(outputStream, charset);
159    return new TeeWriter(writer);
160  }
161}