/* Copyright (c) 2020 Dovecot authors, see the included COPYING file */

#include "lib.h"

#include "lib-signals.h"

#include "time-util.h"

#include "cpu-limit.h"

#include <sys/time.h>

#include <sys/resource.h>

struct cpu_limit {

  struct cpu_limit *parent;

  enum cpu_limit_type type;

  unsigned int cpu_limit_secs;

  struct rusage initial_usage;

  bool limit_reached;

};

static struct cpu_limit *cpu_limit;

static struct rlimit orig_limit, last_set_rlimit;

static volatile sig_atomic_t xcpu_signal_counter;

static sig_atomic_t checked_signal_counter;

static unsigned int rlim_cur_adjust_secs;

static void

cpu_limit_handler(const siginfo_t *si ATTR_UNUSED, void *context ATTR_UNUSED)

{

  xcpu_signal_counter++;

}

static unsigned int

cpu_limit_get_usage_msecs_with(struct cpu_limit *climit,

             enum cpu_limit_type type,

             const struct rusage *rusage)

{

  struct timeval cpu_usage = { 0, 0 };

  long long usage_diff;

  if ((type & CPU_LIMIT_TYPE_USER) != 0)

    timeval_add(&cpu_usage, &rusage->ru_utime);

  if ((type & CPU_LIMIT_TYPE_SYSTEM) != 0)

    timeval_add(&cpu_usage, &rusage->ru_stime);

  struct timeval initial_total = { 0, 0 };

  if ((type & CPU_LIMIT_TYPE_USER) != 0)

    timeval_add(&initial_total, &climit->initial_usage.ru_utime);

  if ((type & CPU_LIMIT_TYPE_SYSTEM) != 0)

    timeval_add(&initial_total, &climit->initial_usage.ru_stime);

  usage_diff = timeval_diff_msecs(&cpu_usage, &initial_total);

  i_assert(usage_diff >= 0);

  i_assert(usage_diff <= UINT_MAX);

  return (unsigned int)usage_diff;

}

unsigned int

cpu_limit_get_usage_msecs(struct cpu_limit *climit, enum cpu_limit_type type)

{

  struct rusage rusage;

  /* Query cpu usage so far */

  if (getrusage(RUSAGE_SELF, &rusage) < 0)

    i_fatal("getrusage() failed: %m");

  return cpu_limit_get_usage_msecs_with(climit, type, &rusage);

}

static bool

cpu_limit_update_recursive(struct cpu_limit *climit,

         const struct rusage *rusage,

         unsigned int *max_wait_secs)

{

  if (climit == NULL)

    return FALSE;

  if (cpu_limit_update_recursive(climit->parent, rusage, max_wait_secs)) {

    /* parent's limit reached */

    climit->limit_reached = TRUE;

    return TRUE;

  }

  unsigned int secs_used =

    cpu_limit_get_usage_msecs_with(climit, climit->type, rusage)/1000;

  if (secs_used >= climit->cpu_limit_secs) {

    climit->limit_reached = TRUE;

    return TRUE;

  }

  unsigned int secs_left = climit->cpu_limit_secs - secs_used;

  if (*max_wait_secs > secs_left)

    *max_wait_secs = secs_left;

  return FALSE;

}

static void cpu_limit_update_rlimit(void)

{

  struct rusage rusage;

  struct rlimit rlimit;

  unsigned int max_wait_secs = UINT_MAX;

  if (getrusage(RUSAGE_SELF, &rusage) < 0)

    i_fatal("getrusage() failed: %m");

  (void)cpu_limit_update_recursive(cpu_limit, &rusage, &max_wait_secs);

  if (max_wait_secs == UINT_MAX) {

    /* All the limits have reached now. Restore the original

       limits. */

    rlimit = orig_limit;

  } else {

    struct timeval tv_limit = rusage.ru_utime;

    timeval_add(&tv_limit, &rusage.ru_stime);

    i_zero(&rlimit);

    /* Add +1 second to round up. */

    rlimit.rlim_cur = tv_limit.tv_sec +

      max_wait_secs + 1 + rlim_cur_adjust_secs;

    rlimit.rlim_max = orig_limit.rlim_max;

  }

  if (last_set_rlimit.rlim_cur != rlimit.rlim_cur) {

    last_set_rlimit = rlimit;

    if (setrlimit(RLIMIT_CPU, &rlimit) < 0)

      i_fatal("setrlimit() failed: %m");

  }

}

struct cpu_limit *

cpu_limit_init(unsigned int cpu_limit_secs, enum cpu_limit_type type)

{

  struct cpu_limit *climit;

  struct rusage rusage;

  i_assert(cpu_limit_secs > 0);

  i_assert(type != 0);

  climit = i_new(struct cpu_limit, 1);

  climit->parent = cpu_limit;

  climit->type = type;

  climit->cpu_limit_secs = cpu_limit_secs;

  /* Query current limit */

  if (climit->parent == NULL) {

    if (getrlimit(RLIMIT_CPU, &orig_limit) < 0)

      i_fatal("getrlimit() failed: %m");

  }

  /* Query cpu usage so far */

  if (getrusage(RUSAGE_SELF, &rusage) < 0)

    i_fatal("getrusage() failed: %m");

  climit->initial_usage = rusage;

  if (climit->parent == NULL) {

    lib_signals_set_handler(SIGXCPU, LIBSIG_FLAG_RESTART,

          cpu_limit_handler, NULL);

  }

  cpu_limit = climit;

  cpu_limit_update_rlimit();

  return climit;

}

void cpu_limit_deinit(struct cpu_limit **_climit)

{

  struct cpu_limit *climit = *_climit;

  *_climit = NULL;

  if (climit == NULL)

    return;

  i_assert(climit == cpu_limit);

  cpu_limit = climit->parent;

  cpu_limit_update_rlimit();

  if (climit->parent == NULL)

    lib_signals_unset_handler(SIGXCPU, cpu_limit_handler, NULL);

  i_free(climit);

}

bool cpu_limit_exceeded(struct cpu_limit *climit)

{

  static struct timeval tv_last = { 0, 0 };

  struct timeval tv_now;

  if (checked_signal_counter != xcpu_signal_counter) {

    i_gettimeofday(&tv_now);

    if (tv_last.tv_sec != 0 &&

        timeval_diff_msecs(&tv_now, &tv_last) < 1000) {

      /* Additional sanity check: We're getting here more

         rapidly than once per second. This isn't expected

         to happen, but at least in theory it could happen

         because rlim_cur isn't clearly calculated from just

         the user+system CPU usage. So in case rlim_cur is

         too low and keeps firing XCPU signal, try to

         increase rlim_cur by 1 second. Eventually it should

         become large enough. */

      rlim_cur_adjust_secs++;

    }

    checked_signal_counter = xcpu_signal_counter;

    cpu_limit_update_rlimit();

  }

  return climit->limit_reached;

}