pagespeed/controller/context_registry.h

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

#ifndef PAGESPEED_CONTROLLER_CONTEXT_REGISTRY_H_
#define PAGESPEED_CONTROLLER_CONTEXT_REGISTRY_H_

#include <unistd.h>

#include <memory>
#include <unordered_set>

#include "base/logging.h"
#include "pagespeed/kernel/base/abstract_mutex.h"
#include "pagespeed/kernel/base/basictypes.h"
#include "pagespeed/kernel/base/condvar.h"
#include "pagespeed/kernel/base/thread_annotations.h"
#include "pagespeed/kernel/base/thread_system.h"

namespace net_instaweb {

// To cleanly shut down gRPC, either client or server side, you need to call
// TryCancel on all outstanding {Client,Server}Contexts and then wait for the
// cancellations to process. ContextRegistry holds a list of active Contexts and
// provides the blocking CancelAllActive() method which cancels all contained
// contexts and then waits for them to be removed from the Registry.

template <typename ContextT>
class ContextRegistry {
 public:
  ContextRegistry(ThreadSystem* thread_system);
  virtual ~ContextRegistry();

  // Returns whether the ContextT was registered or not. Will only fail once
  // CancelAllActive has been called.
  bool TryRegisterContext(ContextT* ctx)
      LOCKS_EXCLUDED(mutex_) WARN_UNUSED_RESULT;
  void RemoveContext(ContextT* ctx) LOCKS_EXCLUDED(mutex_);

  // Calls TryCancel on all contained Contexts and then blocks until all have
  // been Removed. Some other thread(s) must process the cancellations or this
  // will block forever. Note that mutex_ is held while TryCancel is called, so
  // TryCancel must not call back into this registry. gRPC delivers
  // cancellations asynchronously afer TryCancel() has returned, so that is not
  // a problem for the intended use.
  void CancelAllActiveAndWait() LOCKS_EXCLUDED(mutex_);

  // Calls TryCancel on all Contained Contexts and then returns immediately.
  // As above, mutex_ is held while TryCancel is called, so TryCancel must
  // not call back into the registry.
  void CancelAllActive() LOCKS_EXCLUDED(mutex_);

  // Whether CancelAllActive has been called yet. Once this starts returning
  // true, it will never again return false. Because of this, a true return
  // can safely be used to skip work on the assumption that TryRegisterContext
  // will fail. However the converse is NOT true: You must use
  // TryRegisterContext to check if it is safe to to work.
  bool IsShutdown() const LOCKS_EXCLUDED(mutex_);

  // Number of contained contexts.
  int Size() const LOCKS_EXCLUDED(mutex_);

  // For use in tests.
  bool Empty() const LOCKS_EXCLUDED(mutex_);

 private:
  typedef std::unordered_set<ContextT*> ContextSet;

  std::unique_ptr<ThreadSystem::CondvarCapableMutex> mutex_;
  std::unique_ptr<ThreadSystem::Condvar> condvar_ GUARDED_BY(mutex_);
  bool shutdown_ GUARDED_BY(mutex_);
  ContextSet contexts_ GUARDED_BY(mutex_);
};

template <typename ContextT>
ContextRegistry<ContextT>::ContextRegistry(ThreadSystem* thread_system)
    : mutex_(thread_system->NewMutex()),
      condvar_(mutex_->NewCondvar()),
      shutdown_(false) {}

template <typename ContextT>
ContextRegistry<ContextT>::~ContextRegistry() {
  DCHECK_EQ(contexts_.size(), 0u);
}

template <typename ContextT>
bool ContextRegistry<ContextT>::Empty() const {
  ScopedMutex lock(mutex_.get());
  return contexts_.empty();
}

template <typename ContextT>
int ContextRegistry<ContextT>::Size() const {
  ScopedMutex lock(mutex_.get());
  return static_cast<int>(contexts_.size());
}

template <typename ContextT>
bool ContextRegistry<ContextT>::IsShutdown() const {
  ScopedMutex lock(mutex_.get());
  return shutdown_;
}

template <typename ContextT>
bool ContextRegistry<ContextT>::TryRegisterContext(ContextT* context) {
  CHECK(context != nullptr);

  ScopedMutex lock(mutex_.get());
  bool inserted = false;
  if (!shutdown_) {
    inserted = contexts_.insert(context).second;
    DCHECK(inserted);
  }
  return inserted;
}

template <typename ContextT>
void ContextRegistry<ContextT>::RemoveContext(ContextT* context) {
  ScopedMutex lock(mutex_.get());
  int num_erased = contexts_.erase(context);
  DCHECK_EQ(num_erased, 1);
  if (num_erased > 0 && contexts_.empty() && shutdown_) {
    condvar_->Broadcast();
  }
}

template <typename ContextT>
void ContextRegistry<ContextT>::CancelAllActive() {
  ContextSet old_contexts;
  {
    ScopedMutex lock(mutex_.get());
    shutdown_ = true;
    // Copy clients so we can iterate over it without the mutex held.
    // shutdown_ will prevent any additional contexts from being added.
    // This cannot use the usual "swap" trick because we want to wait for the
    // contexts to be removed via calls to RemoveContext on another thread.
    old_contexts = contexts_;
  }

  // If there is nothing to do, we might as well avoid taking the lock a second
  // time, below.
  if (old_contexts.empty()) {
    return;
  }

  for (ContextT* ctx : old_contexts) {
    // There might be a few hundred entries in old_contexts, so we prefer not to
    // hold the lock while we iterate through all of them. However, as soon as
    // we release the lock, contexts that finish naturally can call back into
    // RemoveContext(), after which they are deleted. Thus, we must check every
    // pointer is still in contexts_ (ie: alive) before we call TryCancel() on
    // it.
    mutex_->Lock();
    if (contexts_.find(ctx) != contexts_.end()) {
      // Per gRPC documentation, it's safe to call this no matter where in the
      // lifecycle we are.
      ctx->TryCancel();
    }
    mutex_->Unlock();
    usleep(1);  // yield to other threads.
  }
}

template <typename ContextT>
void ContextRegistry<ContextT>::CancelAllActiveAndWait() {
  CancelAllActive();

  // Now wait for contexts_ to drain as the Cancel events are processed.
  {
    ScopedMutex lock(mutex_.get());
    while (!contexts_.empty()) {
      condvar_->Wait();
    }
  }
}

}  // namespace net_instaweb

#endif  // PAGESPEED_CONTROLLER_CONTEXT_REGISTRY_H_