/* Copyright 2019 The OpenXLA Authors.

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

#ifndef XLA_PJRT_SEMAPHORE_H_
#define XLA_PJRT_SEMAPHORE_H_

#include <cstdint>

#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "xla/types.h"

namespace xla {

class Semaphore {
 public:
  explicit Semaphore(int64_t capacity);

  // Acquires `amount` units. Blocks until `amount` units are available.
  void Acquire(int64_t amount);

  // Tries to acquire `amount` units. Returns true if successful. Returns false
  // immediately if not enough units are available.
  bool TryAcquire(int64_t amount);

  // Returns `amount` units to the semaphore.
  void Release(int64_t amount);

  // Returns the capacity of the semaphore.
  int64_t capacity() const { return max_capacity_; }

  class ScopedReservation {
   public:
    ScopedReservation(Semaphore* semaphore, int64_t amount)
        : semaphore_(semaphore), amount_(amount) {}
    ~ScopedReservation();

    ScopedReservation(const ScopedReservation&) = delete;
    ScopedReservation(ScopedReservation&& other) noexcept;
    ScopedReservation& operator=(const ScopedReservation&) = delete;
    ScopedReservation& operator=(ScopedReservation&& other) noexcept;
    int64_t amount() const { return amount_; }

   private:
    Semaphore* semaphore_;
    int64_t amount_;
  };
  // RAII version of Acquire. Releases the reservation when the
  // ScopedReservation is destroyed.
  ScopedReservation ScopedAcquire(int64_t amount);

 private:
  struct CanAcquireArgs {
    Semaphore* semaphore;
    int64_t amount;
  };
  static bool CanAcquire(CanAcquireArgs* args)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(args->semaphore->mu_);

  absl::Mutex mu_;
  int64_t value_ ABSL_GUARDED_BY(mu_);
  const int64_t max_capacity_;
};

}  // namespace xla

#endif  // XLA_PJRT_SEMAPHORE_H_