package water;
import water.util.ArrayUtils;
import water.util.Log;

import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Stack;

// A "scope" for tracking Key lifetimes.
//
// A Scope defines a *SINGLE THREADED* local lifetime management context,
// stored in Thread Local Storage.  Scopes can be explicitly entered or exited.
// User keys created by this thread are tracked, and deleted when the scope is
// exited.  Since enter & exit are explicit, failure to exit means the Keys
// leak (there is no reliable thread-on-exit cleanup action).  You must call
// Scope.exit() at some point.  Only user keys & Vec keys are tracked.
//
// Scopes support nesting.  Scopes support partial cleanup: you can list Keys
// you'd like to keep in the exit() call.  These will be "bumped up" to the
// higher nested scope - or escaped and become untracked at the top-level.

public class Scope {
  // Thread-based Key lifetime tracking
  static private final ThreadLocal<Scope> _scope = new ThreadLocal<Scope>() {
    @Override protected Scope initialValue() { return new Scope(); }
  };
  private final Stack<HashSet<Key>> _keys = new Stack<>();

  static public void enter() { _scope.get()._keys.push(new HashSet<Key>()); }
  static public void exit () { exit(null); }

  static public Key[] exit(Key... keep) {
    Key[] arrkeep = null;
    if (keep != null) arrkeep = (Key[]) Arrays.asList(keep).toArray();
    Arrays.sort(arrkeep);
    Stack<HashSet<Key>> keys = _scope.get()._keys;
    if (keys.size() > 0) {
      for (Key key : keys.pop()) {
        int found = -1;
        if (arrkeep!=null) found = Arrays.binarySearch(arrkeep, key);
        if (found >= 0) {
          Log.info("Not deleting key " + key);
        }
        if (arrkeep.length == 0 || found < 0) {
          Log.info("Deleting key " + key);
          Keyed.remove(key);
        }
      }
    }
    return keep;
  }

  static public void track( Key k ) {
    if( !k.user_allowed() && !k.isVec() ) return; // Not tracked
    Scope scope = _scope.get();                   // Pay the price of T.L.S. lookup
    if( scope == null ) return; // Not tracking this thread
    if( scope._keys.size() == 0 ) return; // Tracked in the past, but no scope now
    scope._keys.peek().add(k);            // Track key
  }
}

package water;

import java.util.*;

// A "scope" for tracking Key lifetimes.
//
// A Scope defines a *SINGLE THREADED* local lifetime management context,
// stored in Thread Local Storage.  Scopes can be explicitly entered or exited.
// User keys created by this thread are tracked, and deleted when the scope is
// exited.  Since enter & exit are explicit, failure to exit means the Keys
// leak (there is no reliable thread-on-exit cleanup action).  You must call
// Scope.exit() at some point.  Only user keys & Vec keys are tracked.
//
// Scopes support nesting.  Scopes support partial cleanup: you can list Keys
// you'd like to keep in the exit() call.  These will be "bumped up" to the
// higher nested scope - or escaped and become untracked at the top-level.

public class Scope {
  // Thread-based Key lifetime tracking
  static private final ThreadLocal<Scope> _scope = new ThreadLocal<Scope>() {
    @Override protected Scope initialValue() { return new Scope(); }
  };
  private final Stack<HashSet<Key>> _keys = new Stack<>();

  static public void enter() { _scope.get()._keys.push(new HashSet<Key>()); }
  static public void exit () { exit(null); }

  static public Key[] exit(Key... keep) {
    List<Key> keylist = new ArrayList<>();
    if (keep != null) {
      for (Key k : Arrays.asList(keep)) {
        if (k != null)
          keylist.add(k);
      }
    }
    Object[] arrkeep = keylist.toArray();
    Arrays.sort(arrkeep);
    Stack<HashSet<Key>> keys = _scope.get()._keys;
    if (keys.size() > 0) {
      for (Key key : keys.pop()) {
        int found = -1;
        if (arrkeep != null) found = Arrays.binarySearch(arrkeep, key);
//        if (found >= 0) {
//          Log.info("Not deleting key " + key);
//        }
        if (arrkeep.length == 0 || found < 0) {
//          Log.info("Deleting key " + key);
          Keyed.remove(key);
        }
      }
    }
    return keep;
  }

  static public void track( Key k ) {
    if( !k.user_allowed() && !k.isVec() ) return; // Not tracked
    Scope scope = _scope.get();                   // Pay the price of T.L.S. lookup
    if( scope == null ) return; // Not tracking this thread
    if( scope._keys.size() == 0 ) return; // Tracked in the past, but no scope now
    if (!scope._keys.peek().contains(k)) {
//      Log.info("Tracking key: " + k);
      scope._keys.peek().add(k);            // Track key
    }
  }
}

Improve Scope code for exit(Key... key), allow null to be in varargs.