001/**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hdfs.util;
019
020/** 
021 * a class to throttle the data transfers.
022 * This class is thread safe. It can be shared by multiple threads.
023 * The parameter bandwidthPerSec specifies the total bandwidth shared by
024 * threads.
025 */
026public class DataTransferThrottler {
027  private long period;          // period over which bw is imposed
028  private long periodExtension; // Max period over which bw accumulates.
029  private long bytesPerPeriod; // total number of bytes can be sent in each period
030  private long curPeriodStart; // current period starting time
031  private long curReserve;     // remaining bytes can be sent in the period
032  private long bytesAlreadyUsed;
033
034  /** Constructor 
035   * @param bandwidthPerSec bandwidth allowed in bytes per second. 
036   */
037  public DataTransferThrottler(long bandwidthPerSec) {
038    this(500, bandwidthPerSec);  // by default throttling period is 500ms 
039  }
040
041  /**
042   * Constructor
043   * @param period in milliseconds. Bandwidth is enforced over this
044   *        period.
045   * @param bandwidthPerSec bandwidth allowed in bytes per second. 
046   */
047  public DataTransferThrottler(long period, long bandwidthPerSec) {
048    this.curPeriodStart = System.currentTimeMillis();
049    this.period = period;
050    this.curReserve = this.bytesPerPeriod = bandwidthPerSec*period/1000;
051    this.periodExtension = period*3;
052  }
053
054  /**
055   * @return current throttle bandwidth in bytes per second.
056   */
057  public synchronized long getBandwidth() {
058    return bytesPerPeriod*1000/period;
059  }
060  
061  /**
062   * Sets throttle bandwidth. This takes affect latest by the end of current
063   * period.
064   * 
065   * @param bytesPerSecond 
066   */
067  public synchronized void setBandwidth(long bytesPerSecond) {
068    if ( bytesPerSecond <= 0 ) {
069      throw new IllegalArgumentException("" + bytesPerSecond);
070    }
071    bytesPerPeriod = bytesPerSecond*period/1000;
072  }
073  
074  /** Given the numOfBytes sent/received since last time throttle was called,
075   * make the current thread sleep if I/O rate is too fast
076   * compared to the given bandwidth.
077   *
078   * @param numOfBytes
079   *     number of bytes sent/received since last time throttle was called
080   */
081  public synchronized void throttle(long numOfBytes) {
082    if ( numOfBytes <= 0 ) {
083      return;
084    }
085
086    curReserve -= numOfBytes;
087    bytesAlreadyUsed += numOfBytes;
088
089    while (curReserve <= 0) {
090      long now = System.currentTimeMillis();
091      long curPeriodEnd = curPeriodStart + period;
092
093      if ( now < curPeriodEnd ) {
094        // Wait for next period so that curReserve can be increased.
095        try {
096          wait( curPeriodEnd - now );
097        } catch (InterruptedException ignored) {}
098      } else if ( now <  (curPeriodStart + periodExtension)) {
099        curPeriodStart = curPeriodEnd;
100        curReserve += bytesPerPeriod;
101      } else {
102        // discard the prev period. Throttler might not have
103        // been used for a long time.
104        curPeriodStart = now;
105        curReserve = bytesPerPeriod - bytesAlreadyUsed;
106      }
107    }
108
109    bytesAlreadyUsed -= numOfBytes;
110  }
111}