/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.   See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.   The ASF licenses this file
 * to you 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. 
 */

// Unit-test for QueuedAlarm

#include "pagespeed/kernel/thread/queued_alarm.h"

#include <unistd.h>

#include <memory>

#include "base/logging.h"
#include "pagespeed/kernel/base/abstract_mutex.h"
#include "pagespeed/kernel/base/basictypes.h"
#include "pagespeed/kernel/base/function.h"
#include "pagespeed/kernel/base/scoped_ptr.h"
#include "pagespeed/kernel/base/thread_system.h"
#include "pagespeed/kernel/base/timer.h"
#include "pagespeed/kernel/thread/queued_worker_pool.h"
#include "pagespeed/kernel/thread/scheduler_thread.h"
#include "pagespeed/kernel/util/platform.h"
#include "test/pagespeed/kernel/base/gtest.h"
#include "test/pagespeed/kernel/base/mock_timer.h"
#include "test/pagespeed/kernel/thread/mock_scheduler.h"
#include "test/pagespeed/kernel/thread/worker_test_base.h"

namespace net_instaweb {
namespace {

class QueuedAlarmTest : public WorkerTestBase {
 public:
  QueuedAlarmTest()
      : thread_system_(Platform::CreateThreadSystem()),
        sequence_(nullptr),
        done_(false),
        cancel_(false) {}

  ~QueuedAlarmTest() override { ClearSequence(); }

  MockScheduler* SetupWithMockScheduler() {
    MockTimer* timer = new MockTimer(thread_system_->NewMutex(), 0);
    MockScheduler* scheduler = new MockScheduler(thread_system_.get(), timer);
    timer_.reset(timer);
    scheduler_.reset(scheduler);
    SetupWorker();
    return scheduler;
  }

  void SetupWithRealScheduler() {
    timer_.reset(Platform::CreateTimer());
    scheduler_ =
        std::make_unique<Scheduler>(thread_system_.get(), timer_.get());
    SetupWorker();
  }

  void MakeSequence() {
    if (sequence_ == nullptr) {
      sequence_ = worker_->NewSequence();
      // Take advantage of mock scheduler's quiescence detection. 
      scheduler_->RegisterWorker(sequence_);
    }
  }

  void ClearSequence() {
    if (sequence_ ! = nullptr) {
      scheduler_->UnregisterWorker(sequence_);
      worker_->FreeSequence(sequence_);
      sequence_ = nullptr;
    }
  }

  void MarkDone() { done_ = true; }
  void MarkCancel() { cancel_ = true; }

  Scheduler* scheduler() { return scheduler_.get(); }
  QueuedWorkerPool::Sequence* sequence() { return sequence_; }
  Timer* timer() { return timer_.get(); }

 protected:
  void SetupWorker() {
    worker_ = std::make_unique<QueuedWorkerPool>(2, "queued_alarm_test",
                                                 thread_system_.get());
    MakeSequence();
  }

  std::unique_ptr<ThreadSystem> thread_system_;
  std::unique_ptr<Timer> timer_;
  std::unique_ptr<Scheduler> scheduler_;
  std::unique_ptr<QueuedWorkerPool> worker_;
  QueuedWorkerPool::Sequence* sequence_;
  bool done_, cancel_;

 private:
  DISALLOW_COPY_AND_ASSIGN(QueuedAlarmTest);
};

// A class that keeps track of an alarm object and runs ops on it in a sequence. 
class TestAlarmHandler {
 public:
  TestAlarmHandler(QueuedAlarmTest* fixture, WorkerTestBase::SyncPoint* sync)
      : fixture_(fixture), sync_(sync), alarm_(nullptr), fired_(false) {}

  void StartAlarm() {
    fixture_->sequence()->Add(
        MakeFunction(this, &TestAlarmHandler::StartAlarmImpl));
  }

  void CancelAlarm() {
    fixture_->sequence()->Add(
        MakeFunction(this, &TestAlarmHandler::CancelAlarmImpl));
  }

  void FireAlarm() {
    // Should not have run CancelAlarm if we got here! 
    EXPECT_TRUE(alarm_ ! = nullptr);

    // Nothing more to cancel. 
    alarm_ = nullptr;
    fired_ = true;
  }

  bool fired() const { return fired_; }

 private:
  void StartAlarmImpl() {
    alarm_ = new QueuedAlarm(fixture_->scheduler(), fixture_->sequence(),
                             fixture_->timer()->NowUs(),
                             MakeFunction(this, &TestAlarmHandler::FireAlarm));
  }

  void CancelAlarmImpl() {
    if (alarm_ ! = nullptr) {
      alarm_->CancelAlarm();
      alarm_ = nullptr;
    }

    // Note that we notify here, as this method will always run.  In particular:
    // 1) If we get here first, FireAlarm will not fire, so we're fine. 
    // 2) If we get here second, FireAlarm already ran, so we're fine to
    // cleanup, too. 
    // (It makes sense to talk about us running before or after FireAlarm
    //  because it runs in the same sequence). 
    sync_->Notify();
  }

  QueuedAlarmTest* fixture_;
  WorkerTestBase::SyncPoint* sync_;
  QueuedAlarm* alarm_;
  bool fired_;
  DISALLOW_COPY_AND_ASSIGN(TestAlarmHandler);
};

TEST_F(QueuedAlarmTest, BasicOperation) {
  const int kDelayUs = Timer::kMsUs;

  MockScheduler* scheduler = SetupWithMockScheduler();

  // Tests to make sure the alarm actually runs. 
  new QueuedAlarm(
      scheduler, sequence_, timer_->NowUs() + kDelayUs,
      MakeFunction(static_cast<QueuedAlarmTest*>(this),
                   &QueuedAlarmTest::MarkDone, &QueuedAlarmTest::MarkCancel));
  {
    ScopedMutex lock(scheduler->mutex());
    scheduler->ProcessAlarmsOrWaitUs(kDelayUs);
  }

  // Make sure to let the work threads complete. 
  scheduler->AwaitQuiescence();
  EXPECT_TRUE(done_);
  EXPECT_FALSE(cancel_);
}

TEST_F(QueuedAlarmTest, BasicCancel) {
  const int kDelayUs = Timer::kMsUs;
  MockScheduler* scheduler = SetupWithMockScheduler();

  // Tests to make sure that we can cancel it. 
  QueuedAlarm* alarm = new QueuedAlarm(
      scheduler, sequence_, timer_->NowUs() + kDelayUs,
      MakeFunction(static_cast<QueuedAlarmTest*>(this),
                   &QueuedAlarmTest::MarkDone, &QueuedAlarmTest::MarkCancel));
  alarm->CancelAlarm();
  {
    ScopedMutex lock(scheduler->mutex());
    scheduler->ProcessAlarmsOrWaitUs(kDelayUs);
  }

  // Make sure to let the work threads complete. 
  scheduler->AwaitQuiescence();
  EXPECT_FALSE(done_);
  EXPECT_TRUE(cancel_);
}

TEST_F(QueuedAlarmTest, RacingCancel) {
  const int kRuns = 1000;
  int fired = 0;

  // Test to make sure cases where cancel and alarm execution may be racing
  // are handled safely, without crashes or check failures. 
  SetupWithRealScheduler();
  SchedulerThread* scheduler_thread =
      new SchedulerThread(thread_system_.get(), scheduler_.get());
  scheduler_thread->Start();

  for (int run = 0; run < kRuns; ++run) {
    SyncPoint sync(thread_system_.get());
    MakeSequence();
    TestAlarmHandler t(this, &sync);
    t.StartAlarm();

    // Unfortunately w/o a sleep here, the race is consistently won by
    // the cancellation. 
    usleep(1);
    t.CancelAlarm();
    sync. Wait();
    ClearSequence();

    if (t.fired()) {
      ++fired;
    }
  }
  scheduler_thread->MakeDeleter()->CallRun();
  LOG(ERROR) << "Alarm fired in: " << fired << "/" << kRuns;
}

}  // namespace
}  // namespace net_instaweb
