YugabyteDB (2.13.1.0-b60, 21121d69985fbf76aa6958d8f04a9bfa936293b5)

//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.

//  This source code is licensed under the BSD-style license found in the

//  LICENSE file in the root directory of this source tree. An additional grant

//  of patent rights can be found in the PATENTS file in the same directory.

//

// The following only applies to changes made to this file as part of YugaByte development.

//

// Portions Copyright (c) YugaByte, Inc.

//

// 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.

//

#pragma once

#include <memory>

namespace rocksdb {

class Env;

class WriteControllerToken;

// WriteController is controlling write stalls in our write code-path. Write

// stalls happen when compaction can't keep up with write rate.

// All of the methods here (including WriteControllerToken's destructors) need

// to be called while holding DB mutex

class WriteController {

 public:

  explicit WriteController(uint64_t _delayed_write_rate = 1024u * 1024u * 32u)

      : total_stopped_(0),

        total_delayed_(0),

        total_compaction_pressure_(0),

        bytes_left_(0),

        last_refill_time_(0) {

    set_delayed_write_rate(_delayed_write_rate);

  }

  ~WriteController() = default;

  // When an actor (column family) requests a stop token, all writes will be

  // stopped until the stop token is released (deleted)

  std::unique_ptr<WriteControllerToken> GetStopToken();

  // When an actor (column family) requests a delay token, total delay for all

  // writes to the DB will be controlled under the delayed write rate. Every

  // write needs to call GetDelay() with number of bytes writing to the DB,

  // which returns number of microseconds to sleep.

  std::unique_ptr<WriteControllerToken> GetDelayToken(

      uint64_t delayed_write_rate);

  // When an actor (column family) requests a moderate token, compaction

  // threads will be increased

  std::unique_ptr<WriteControllerToken> GetCompactionPressureToken();

  // these three metods are querying the state of the WriteController

  bool IsStopped() const;

  bool NeedsDelay() const { return total_delayed_ > 0; }

  bool NeedSpeedupCompaction() const {

    return IsStopped() || 
NeedsDelay()1.25M
 || 
total_compaction_pressure_ > 01.25M
;

  }

  // return how many microseconds the caller needs to sleep after the call

  // num_bytes: how many number of bytes to put into the DB.

  // Prerequisite: DB mutex held.

  uint64_t GetDelay(Env* env, uint64_t num_bytes);

  void set_delayed_write_rate(uint64_t write_rate) {

    // avoid divide 0

    if (write_rate == 0) {

      write_rate = 1u;

    }

    delayed_write_rate_ = write_rate;

  }

  uint64_t delayed_write_rate() const { return delayed_write_rate_; }

 private:

  friend class WriteControllerToken;

  friend class StopWriteToken;

  friend class DelayWriteToken;

  friend class CompactionPressureToken;

  int total_stopped_;

  int total_delayed_;

  int total_compaction_pressure_;

  uint64_t bytes_left_;

  uint64_t last_refill_time_;

  uint64_t delayed_write_rate_;

};

class WriteControllerToken {

 public:

  explicit WriteControllerToken(WriteController* controller)

      : controller_(controller) {}

  virtual ~WriteControllerToken() {}

 protected:

  WriteController* controller_;

 private:

  // no copying allowed

  WriteControllerToken(const WriteControllerToken&) = delete;

  void operator=(const WriteControllerToken&) = delete;

};

class StopWriteToken : public WriteControllerToken {

 public:

  explicit StopWriteToken(WriteController* controller)

      : WriteControllerToken(controller) {}

  virtual ~StopWriteToken();

};

class DelayWriteToken : public WriteControllerToken {

 public:

  explicit DelayWriteToken(WriteController* controller)

      : WriteControllerToken(controller) {}

  virtual ~DelayWriteToken();

};

class CompactionPressureToken : public WriteControllerToken {

 public:

  explicit CompactionPressureToken(WriteController* controller)

      : WriteControllerToken(controller) {}

  virtual ~CompactionPressureToken();

};

}  // namespace rocksdb