<commit_before>package water;

import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryNotificationInfo;
import java.lang.management.MemoryPoolMXBean;
import java.lang.management.MemoryType;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicLong;

import javax.management.Notification;
import javax.management.NotificationEmitter;

/**
 * Manages memory assigned to key/value pairs. All byte arrays used in
 * keys/values should be allocated through this class - otherwise we risking
 * running out of java memory, and throw unexpected OutOfMemory errors. The
 * theory here is that *most* allocated bytes are allocated in large chunks by
 * allocating new Values - with large backing arrays. If we intercept these
 * allocation points, we cover most Java allocations. If such an allocation
 * might trigger an OOM error we first free up some other memory.
 *
 * MemoryManager monitors memory used by the K/V store (by walking through the
 * store (see Cleaner) and overall heap usage by hooking into gc.
 *
 * Memory is freed if either the cached memory is above the limit or if the
 * overall heap usage is too high (in which case we want to use less mem for
 * cache). There is also a lower limit on the amount of cache so that we never
 * delete all the cache and therefore some computation should always be able to
 * progress.
 *
 * The amount of memory to be freed is determined as the max of cached mem above
 * the limit and heap usage above the limit.
 *
 * One of the primary control inputs is FullGC cycles: we check heap usage and
 * set guidance for cache levels. We assume after a FullGC that the heap only
 * has POJOs (Plain Old Java Objects, unknown size) and K/V Cached stuff
 * (counted by us). We compute the free heap as MEM_MAX-heapUsage (after GC),
 * and we compute POJO size as (heapUsage - K/V cache usage).
 *
 * @author tomas
 * @author cliffc
 */
public abstract class MemoryManager {

  // max heap memory
  static final long MEM_MAX = Runtime.getRuntime().maxMemory();

  // Memory currently being used by running tasks
  static final AtomicLong _taskMem = new AtomicLong(0);

  /**
   * Try to reserve memory needed for task execution and return true if succeeded.
   * Tasks have a shared pool of memory which they should ask for in advance before they even try to allocate it.
   * We assume here, that tasks can take all the memory (except 1/8th of the heap which is set as minimum cache size).
   *
   * This method is another backpressure mechanism to make sure we do not exhaust system's resources by running too many tasks at the same time.
   * Tasks are expected to reserve memory before proceeding with their execution and making sure they release it when done.
   *
   * @param m - requested number of bytes
   * @return true if there is enough free memory
   */
  public static boolean tryReserveTaskMem(long m){
    long current = _taskMem.addAndGet(m);
    if(current > (MEM_MAX - (MEM_MAX >> 3))){ // 1/8th of the heap is reserved for cache
      current = _taskMem.addAndGet(-m);
      return false;
    }
    return true;
  }

  /**
   * Free the memory successfully reserved by task.
   * @param m
   */
  public static void freeTaskMem(long m){
    _taskMem.addAndGet(-m);
  }

  // Callbacks from GC
  static final HeapUsageMonitor HEAP_USAGE_MONITOR = new HeapUsageMonitor();
  static volatile long HEAP_USED_AT_LAST_GC;
  static volatile long POJO_USED_AT_LAST_GC;
  static volatile long TIME_AT_LAST_GC=System.currentTimeMillis();

  // Keep the K/V store below this threshold AND this is the FullGC call-back
  // threshold - which is limited in size to the old-gen pool size.
  static long MEM_CRITICAL = HEAP_USAGE_MONITOR._gc_callback;

  // Block allocations?
  private static volatile boolean CAN_ALLOC = true;

  // Lock for blocking on allocations
  private static Object _lock = new Object();

  // My Histogram. Called from any thread calling into the MM.
  // Singleton, allocated now so I do not allocate during an OOM event.
  static private final H2O.Cleaner.Histo myHisto = new H2O.Cleaner.Histo();


  public static void setMemGood() {
    if( CAN_ALLOC ) return;
    synchronized(_lock) {
      CAN_ALLOC = true;
      System.err.println("Continuing after swapping");
      _lock.notifyAll();
    }
  }
  public static void setMemLow() {
    if( !CAN_ALLOC ) return;
    synchronized(_lock) {
      CAN_ALLOC = false;
      System.err.println("Pausing to swap to disk; more memory may help");
    }
  }

  public static void set_goals( String msg, boolean oom){
    set_goals(msg, oom, 0);
  }
  // Set K/V cache goals.
  // Allow (or disallow) allocations.
  // Called from the Cleaner, when "cacheUsed" has changed significantly.
  // Called from any FullGC notification, and HEAP/POJO_USED changed.
  // Called on any OOM allocation
  public static void set_goals( String msg, boolean oom , long bytes) {
    // Our best guess of free memory, as of the last GC cycle
    long freeHeap = MEM_MAX - HEAP_USED_AT_LAST_GC;
    assert freeHeap >= 0 : "I am really confused about the heap usage";
    // Current memory held in the K/V store.
    long cacheUsage = myHisto.histo(false)._cached;

    // Block allocations if:
    // the cache is > 7/8 MEM_MAX, OR
    // we cannot allocate an equal amount of POJOs, POJO_USED_AT_LAST_GC > freeHeap.
    // Decay POJOS_USED by 1/8th every 5 sec: assume we got hit with a single
    // large allocation which is not repeating - so we do not need to have
    // double the POJO amount.
    // Keep at least 1/8th heap for caching.
    // Emergency-clean the cache down to the blocking level.
    long d = MEM_CRITICAL;
    // Decay POJO amount
    long p = POJO_USED_AT_LAST_GC;
    long age = (System.currentTimeMillis() - TIME_AT_LAST_GC); // Age since last FullGC
    age = Math.min(age,10*60*1000 ); // Clip at 10mins
    while( (age-=5000) > 0 ) p = p-(p>>3); // Decay effective POJO by 1/8th every 5sec
    d -= 2*p - bytes; // Allow for the effective POJO, and again to throttle GC rate
    d = Math.max(d,MEM_MAX>>3); // Keep at least 1/8th heap
    H2O.Cleaner.DESIRED = d;

    if( cacheUsage > H2O.Cleaner.DESIRED ) {
      if( H2O.Cleaner.VERBOSE ) {
        System.out.print(CAN_ALLOC?"Blocking! ":"blocked: ");
        System.out.println(msg+", KV="+(cacheUsage>>20)+"M"+
                           ", POJO="+(POJO_USED_AT_LAST_GC>>20)+"M"+
                           ", free="+(freeHeap>>20)+"M"+
                           ", MAX="+(MEM_MAX>>20)+"M"+
                           ", DESIRED="+(H2O.Cleaner.DESIRED>>20)+"M");
      }
      setMemLow(); // Stop allocations; trigger emergency clean
      H2O.kick_store_cleaner();
    } else { // Else we are not *emergency* cleaning, but may be lazily cleaning.
      if( H2O.Cleaner.VERBOSE && !CAN_ALLOC )
        System.out.println("Unblocking: "+msg+", KV="+(cacheUsage>>20)+"M"+
                           ", POJO="+(POJO_USED_AT_LAST_GC>>20)+"M"+
                           ", free="+(freeHeap>>20)+"M"+
                           ", MAX="+(MEM_MAX>>20)+"M"+
                           ", DESIRED="+(H2O.Cleaner.DESIRED>>20)+"M");
      setMemGood();
      assert !oom:"MEM_MAX = " + MEM_MAX + ", DESIRED = " + d +", CACHE = " + cacheUsage + ", p = " + p + ", bytes = " + bytes; // Confused? OOM should have FullGCd should have set low-mem goals
    }
  }

  public static long reserveTaskMemory(long sz){
    return 0;
  }
  /**
   * Monitors the heap usage after full gc run and tells Cleaner to free memory
   * if mem usage is too high. Stops new allocation if mem usage is critical.
   * @author tomas
   */
  private static class HeapUsageMonitor implements javax.management.NotificationListener {
    MemoryMXBean _allMemBean = ManagementFactory.getMemoryMXBean(); // general
    MemoryPoolMXBean _oldGenBean;
    public long _gc_callback;<commit_msg>Only block allocation if we throw an actual OOM, even if the cache/heap is very full, and the disk is full so we cannot swap.<commit_after>package water;

import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryNotificationInfo;
import java.lang.management.MemoryPoolMXBean;
import java.lang.management.MemoryType;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicLong;

import javax.management.Notification;
import javax.management.NotificationEmitter;

/**
 * Manages memory assigned to key/value pairs. All byte arrays used in
 * keys/values should be allocated through this class - otherwise we risking
 * running out of java memory, and throw unexpected OutOfMemory errors. The
 * theory here is that *most* allocated bytes are allocated in large chunks by
 * allocating new Values - with large backing arrays. If we intercept these
 * allocation points, we cover most Java allocations. If such an allocation
 * might trigger an OOM error we first free up some other memory.
 *
 * MemoryManager monitors memory used by the K/V store (by walking through the
 * store (see Cleaner) and overall heap usage by hooking into gc.
 *
 * Memory is freed if either the cached memory is above the limit or if the
 * overall heap usage is too high (in which case we want to use less mem for
 * cache). There is also a lower limit on the amount of cache so that we never
 * delete all the cache and therefore some computation should always be able to
 * progress.
 *
 * The amount of memory to be freed is determined as the max of cached mem above
 * the limit and heap usage above the limit.
 *
 * One of the primary control inputs is FullGC cycles: we check heap usage and
 * set guidance for cache levels. We assume after a FullGC that the heap only
 * has POJOs (Plain Old Java Objects, unknown size) and K/V Cached stuff
 * (counted by us). We compute the free heap as MEM_MAX-heapUsage (after GC),
 * and we compute POJO size as (heapUsage - K/V cache usage).
 *
 * @author tomas
 * @author cliffc
 */
public abstract class MemoryManager {

  // max heap memory
  static final long MEM_MAX = Runtime.getRuntime().maxMemory();

  // Memory currently being used by running tasks
  static final AtomicLong _taskMem = new AtomicLong(0);

  /**
   * Try to reserve memory needed for task execution and return true if succeeded.
   * Tasks have a shared pool of memory which they should ask for in advance before they even try to allocate it.
   * We assume here, that tasks can take all the memory (except 1/8th of the heap which is set as minimum cache size).
   *
   * This method is another backpressure mechanism to make sure we do not exhaust system's resources by running too many tasks at the same time.
   * Tasks are expected to reserve memory before proceeding with their execution and making sure they release it when done.
   *
   * @param m - requested number of bytes
   * @return true if there is enough free memory
   */
  public static boolean tryReserveTaskMem(long m){
    long current = _taskMem.addAndGet(m);
    if(current > (MEM_MAX - (MEM_MAX >> 3))){ // 1/8th of the heap is reserved for cache
      current = _taskMem.addAndGet(-m);
      return false;
    }
    return true;
  }

  /**
   * Free the memory successfully reserved by task.
   * @param m
   */
  public static void freeTaskMem(long m){
    _taskMem.addAndGet(-m);
  }

  // Callbacks from GC
  static final HeapUsageMonitor HEAP_USAGE_MONITOR = new HeapUsageMonitor();
  static volatile long HEAP_USED_AT_LAST_GC;
  static volatile long POJO_USED_AT_LAST_GC;
  static volatile long TIME_AT_LAST_GC=System.currentTimeMillis();

  // Keep the K/V store below this threshold AND this is the FullGC call-back
  // threshold - which is limited in size to the old-gen pool size.
  static long MEM_CRITICAL = HEAP_USAGE_MONITOR._gc_callback;

  // Block allocations?
  private static volatile boolean CAN_ALLOC = true;

  // Lock for blocking on allocations
  private static Object _lock = new Object();

  // My Histogram. Called from any thread calling into the MM.
  // Singleton, allocated now so I do not allocate during an OOM event.
  static private final H2O.Cleaner.Histo myHisto = new H2O.Cleaner.Histo();


  public static void setMemGood() {
    if( CAN_ALLOC ) return;
    synchronized(_lock) {
      CAN_ALLOC = true;
      System.err.println("Continuing after swapping");
      _lock.notifyAll();
    }
  }
  public static void setMemLow() {
    if( !CAN_ALLOC ) return;
    synchronized(_lock) {
      CAN_ALLOC = false;
      System.err.println("Pausing to swap to disk; more memory may help");
    }
  }

  public static void set_goals( String msg, boolean oom){
    set_goals(msg, oom, 0);
  }
  // Set K/V cache goals.
  // Allow (or disallow) allocations.
  // Called from the Cleaner, when "cacheUsed" has changed significantly.
  // Called from any FullGC notification, and HEAP/POJO_USED changed.
  // Called on any OOM allocation
  public static void set_goals( String msg, boolean oom , long bytes) {
    // Our best guess of free memory, as of the last GC cycle
    long freeHeap = MEM_MAX - HEAP_USED_AT_LAST_GC;
    assert freeHeap >= 0 : "I am really confused about the heap usage";
    // Current memory held in the K/V store.
    long cacheUsage = myHisto.histo(false)._cached;

    // Block allocations if:
    // the cache is > 7/8 MEM_MAX, OR
    // we cannot allocate an equal amount of POJOs, POJO_USED_AT_LAST_GC > freeHeap.
    // Decay POJOS_USED by 1/8th every 5 sec: assume we got hit with a single
    // large allocation which is not repeating - so we do not need to have
    // double the POJO amount.
    // Keep at least 1/8th heap for caching.
    // Emergency-clean the cache down to the blocking level.
    long d = MEM_CRITICAL;
    // Decay POJO amount
    long p = POJO_USED_AT_LAST_GC;
    long age = (System.currentTimeMillis() - TIME_AT_LAST_GC); // Age since last FullGC
    age = Math.min(age,10*60*1000 ); // Clip at 10mins
    while( (age-=5000) > 0 ) p = p-(p>>3); // Decay effective POJO by 1/8th every 5sec
    d -= 2*p - bytes; // Allow for the effective POJO, and again to throttle GC rate
    d = Math.max(d,MEM_MAX>>3); // Keep at least 1/8th heap
    H2O.Cleaner.DESIRED = d;

    if( cacheUsage > H2O.Cleaner.DESIRED ) {
      if( H2O.Cleaner.VERBOSE ) {
        System.out.print(CAN_ALLOC?"Blocking! ":"blocked: ");
        System.out.println(msg+", KV="+(cacheUsage>>20)+"M"+
                           ", POJO="+(POJO_USED_AT_LAST_GC>>20)+"M"+
                           ", free="+(freeHeap>>20)+"M"+
                           ", MAX="+(MEM_MAX>>20)+"M"+
                           ", DESIRED="+(H2O.Cleaner.DESIRED>>20)+"M");
      }
      if( oom ) setMemLow(); // Stop allocations; trigger emergency clean
      H2O.kick_store_cleaner();
    } else { // Else we are not *emergency* cleaning, but may be lazily cleaning.
      if( H2O.Cleaner.VERBOSE && !CAN_ALLOC )
        System.out.println("Unblocking: "+msg+", KV="+(cacheUsage>>20)+"M"+
                           ", POJO="+(POJO_USED_AT_LAST_GC>>20)+"M"+
                           ", free="+(freeHeap>>20)+"M"+
                           ", MAX="+(MEM_MAX>>20)+"M"+
                           ", DESIRED="+(H2O.Cleaner.DESIRED>>20)+"M");
      setMemGood();
      assert !oom:"MEM_MAX = " + MEM_MAX + ", DESIRED = " + d +", CACHE = " + cacheUsage + ", p = " + p + ", bytes = " + bytes; // Confused? OOM should have FullGCd should have set low-mem goals
    }
  }

  public static long reserveTaskMemory(long sz){
    return 0;
  }
  /**
   * Monitors the heap usage after full gc run and tells Cleaner to free memory
   * if mem usage is too high. Stops new allocation if mem usage is critical.
   * @author tomas
   */
  private static class HeapUsageMonitor implements javax.management.NotificationListener {
    MemoryMXBean _allMemBean = ManagementFactory.getMemoryMXBean(); // general
    MemoryPoolMXBean _oldGenBean;
    public long _gc_callback;
{'fraction_non_alphanumeric': 0.09153518657650889, 'fraction_numerical': 0.008514901076884547, 'mean_word_length': 3.00125250501002, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 67, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5202090', 'n_tokens_mistral': 4799, 'n_tokens_neox': 4645, 'n_words': 1844}