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