com.unboundid.util.RateLimitedOutputStream Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of unboundid-ldapsdk Show documentation
The UnboundID LDAP SDK for Java is a fast, comprehensive, and easy-to-use Java API for communicating with LDAP directory servers and performing related tasks like reading and writing LDIF, encoding and decoding data using base64 and ASN.1 BER, and performing secure communication. This package contains the Standard Edition of the LDAP SDK, which is a complete, general-purpose library for communicating with LDAPv3 directory servers.
The newest version!
/*
 * Copyright 2018-2024 Ping Identity Corporation
 * All Rights Reserved.
 */
/*
 * Copyright 2018-2024 Ping Identity Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
/*
 * Copyright (C) 2018-2024 Ping Identity Corporation
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License (GPLv2 only)
 * or the terms of the GNU Lesser General Public License (LGPLv2.1 only)
 * as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see .
 */
package com.unboundid.util;



import java.io.IOException;
import java.io.OutputStream;



/**
 * This class provides an {@code OutputStream} implementation that uses a
 * {@link FixedRateBarrier} to impose an upper bound on the rate (in bytes per
 * second) at which data can be written to a wrapped {@code OutputStream}.
 */
@ThreadSafety(level=ThreadSafetyLevel.NOT_THREADSAFE)
public final class RateLimitedOutputStream
       extends OutputStream
{
  // Indicates whether to automatically flush the stream after each write.
  private final boolean autoFlush;

  // The fixed-rate barrier that will serve as a rate limiter for this class.
  @NotNull private final FixedRateBarrier rateLimiter;

  // The output stream to which the data will actually be written.
  @NotNull private final OutputStream wrappedStream;

  // The maximum number of bytes that can be written in any single call to the
  // rate limiter.
  private final int maxBytesPerWrite;



  /**
   * Creates a new instance of this rate-limited output stream that wraps the
   * provided output stream.
   *
   * @param  wrappedStream      The output stream to which the data will
   *                            actually be written.  It must not be
   *                            {@code null}.
   * @param  maxBytesPerSecond  The maximum number of bytes per second that can
   *                            be written using this output stream.  It must be
   *                            greater than zero.
   * @param  autoFlush          Indicates whether to automatically flush the
   *                            wrapped output stream after each write.
   */
  public RateLimitedOutputStream(@NotNull final OutputStream wrappedStream,
                                 final int maxBytesPerSecond,
                                 final boolean autoFlush)
  {
    Validator.ensureTrue((wrappedStream != null),
         "RateLimitedOutputStream.wrappedStream must not be null.");
    Validator.ensureTrue((maxBytesPerSecond > 0),
         "RateLimitedOutputStream.maxBytesPerSecond must be greater than " +
              "zero.  The provided value was " + maxBytesPerSecond);

    this.wrappedStream = wrappedStream;
    this.autoFlush = autoFlush;

    rateLimiter = new FixedRateBarrier(1000L, maxBytesPerSecond);
    maxBytesPerWrite = Math.max(1, (maxBytesPerSecond / 100));
  }



  /**
   * Closes this output stream and the wrapped stream.
   *
   * @throws  IOException  If a problem is encountered while closing the wrapped
   *                       output stream.
   */
  @Override()
  public void close()
         throws IOException
  {
    wrappedStream.close();
  }



  /**
   * Writes a single byte of data to the wrapped output stream.
   *
   * @param  b  The byte of data to be written.  Only the least significant
   *            eight bits will be written.
   *
   * @throws  IOException  If a problem is encountered while writing to the
   *                       wrapped stream.
   */
  @Override()
  public void write(final int b)
         throws IOException
  {
    rateLimiter.await();
    wrappedStream.write(b);

    if (autoFlush)
    {
      wrappedStream.flush();
    }
  }



  /**
   * Writes the contents of the provided array to the wrapped output stream.
   *
   * @param  b  The byte array containing the data to be written.  It must not
   *            be {@code null}.
   *
   * @throws  IOException  If a problem is encountered while writing to the
   *                       wrapped stream.
   */
  @Override()
  public void write(@NotNull final byte[] b)
         throws IOException
  {
    write(b, 0, b.length);
  }



  /**
   * Writes the contents of the specified portion of the provided array to the
   * wrapped output stream.
   *
   * @param  b       The byte array containing the data to be written.  It must
   *                 not be {@code null}.
   * @param  offset  The position in the provided array at which the data to
   *                 write begins.  It must be greater than or equal to zero and
   *                 less than the length of the provided array.
   * @param  length  The number of bytes to be written.  It must not be
   *                 negative, and the sum of offset and length must be less
   *                 than or equal to the length of the provided array.
   *
   * @throws  IOException  If a problem is encountered while writing to the
   *                       wrapped stream.
   */
  @Override()
  public void write(@NotNull final byte[] b, final int offset, final int length)
         throws IOException
  {
    if (length <= 0)
    {
      return;
    }

    if (length <= maxBytesPerWrite)
    {
      rateLimiter.await(length);
      wrappedStream.write(b, offset, length);
    }
    else
    {
      int pos = offset;
      int remainingToWrite = length;
      while (remainingToWrite > 0)
      {
        final int lengthThisWrite =
             Math.min(remainingToWrite, maxBytesPerWrite);
        rateLimiter.await(lengthThisWrite);
        wrappedStream.write(b, pos, lengthThisWrite);
        pos += lengthThisWrite;
        remainingToWrite -= lengthThisWrite;
      }
    }

    if (autoFlush)
    {
      wrappedStream.flush();
    }
  }



  /**
   * Flushes the contents of the wrapped stream.
   *
   * @throws  IOException  If a problem is encountered while flushing the
   *                       wrapped stream.
   */
  @Override()
  public void flush()
         throws IOException
  {
    wrappedStream.flush();
  }
}