/*
 * Copyright (c) 2023, PostgreSQL Global Development Group
 * See the LICENSE file in the project root for more information.
 */

/* changes were made to move it into the org.postgresql.util package
 *
 * Copyright 2022 Juan Lopes
 *
 * Licensed 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.
 */

package org.postgresql.util;

// import org.checkerframework.checker.nullness.qual.Nullable;

import java.lang.ref.PhantomReference;
import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.time.Duration;
import java.util.concurrent.ForkJoinPool;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * LazyCleaner is a utility class that allows to register objects for deferred cleanup.
 *
 * <p>This is the Java 8 compatible implementation that uses PhantomReferences
 * and ForkJoinPool for deferred cleanup operations.</p>
 *
 * <p>On Java 11+, this class is replaced by a version that uses the native
 * {@link java.lang.ref.Cleaner} API via the multi-release JAR mechanism.</p>
 *
 * <p>Note: this is a driver-internal class</p>
 */
public class LazyCleanerImpl implements LazyCleaner {
  private static final Logger LOGGER = Logger.getLogger(LazyCleanerImpl.class.getName());
  private static final LazyCleanerImpl instance = new LazyCleanerImpl(
      "PostgreSQL-JDBC-Cleaner",
      Duration.ofMillis(Long.getLong("pgjdbc.config.cleanup.thread.ttl", 30000))
  );

  private final ReferenceQueue<Object> queue = new ReferenceQueue<>();
  private final String threadName;
  private final Duration threadTtl;
  private boolean threadRunning;
  private /* @Nullable */ Node<?> first;

  /**
   * Creates a LazyCleaner with the specified configuration.
   *
   * @param threadName the name for the cleanup thread
   * @param threadTtl the maximum time the cleanup thread will wait for references
   */
  public LazyCleanerImpl(String threadName, Duration threadTtl) {
    this.threadName = threadName;
    this.threadTtl = threadTtl;
  }

  /**
   * Returns a default cleaner instance.
   *
   * <p>Note: this is driver-internal API.</p>
   * @return the instance of LazyCleaner
   */
  public static LazyCleanerImpl getInstance() {
    return instance;
  }

  @Override
  public <T extends Throwable> Cleanable<T> register(Object obj, CleaningAction<T> action) {
    assert obj != action : "object handle should not be the same as cleaning action, otherwise"
        + " the object will never become phantom reachable, so the action will never trigger";
    return add(new Node<>(obj, action));
  }

  /**
   * Returns whether the cleanup thread is currently running.
   *
   * @return true if cleanup operations are active
   */
  public synchronized boolean isThreadRunning() {
    return threadRunning;
  }

  private synchronized boolean checkEmpty() {
    if (first == null) {
      threadRunning = false;
      return true;
    }
    return false;
  }

  private synchronized <T extends Throwable> Node<T> add(Node<T> node) {
    if (first != null) {
      node.next = first;
      first.prev = node;
    }
    first = node;

    if (!threadRunning) {
      threadRunning = startThread();
    }
    return node;
  }

  /**
   * RefQueueBlocker retrieves references from the reference queue without blocking ForkJoinPool
   * CPU threads.
   *
   * @param <T> the type of the objects referenced by the {@link Reference}s in the queue
   */
  private static class RefQueueBlocker<T> implements ForkJoinPool.ManagedBlocker {
    private final ReferenceQueue<T> queue;
    private final String threadName;
    private /* @Nullable */ Reference<? extends T> ref;
    private final long blockTimeoutMillis;

    RefQueueBlocker(ReferenceQueue<T> queue, String threadName, Duration blockTimeout) {
      this.queue = queue;
      this.threadName = threadName;
      // ReferenceQueue.remove(0) blocks indefinitely; clamp to a small positive timeout so a
      // user-configured threadTtl of 0/negative still yields a quickly-terminating cleanup task
      // rather than a thread that parks forever waiting for the next reference.
      long millis = blockTimeout.toMillis();
      this.blockTimeoutMillis = millis > 0 ? millis : 1;
    }

    @Override
    public boolean isReleasable() {
      if (ref != null) {
        return true; // already have a ref from a previous call
      }
      // non-blocking check
      ref = queue.poll();
      return ref != null;
    }

    @Override
    public boolean block() throws InterruptedException {
      if (isReleasable()) {
        return true;
      }
      Thread currentThread = Thread.currentThread();
      // ForkJoinPool reuses threads, so we set the thread name just for the blocking operation
      String oldName = currentThread.getName();
      try {
        currentThread.setName(threadName);
        // Wait the full blockTimeoutMillis for a reference. Returning true unconditionally lets
        // managedBlock exit so the outer loop can decide whether to terminate the cleanup task.
        ref = queue.remove(blockTimeoutMillis);
      } finally {
        currentThread.setName(oldName);
      }
      return true;
    }

    /* @Nullable */ Reference<? extends T> drainOne() {
      Reference<? extends T> ref = this.ref;
      this.ref = null;
      return ref;
    }
  }

  private boolean startThread() {
    // We use ForkJoinPool to work around Thread.inheritedAccessControlContext memory leak
    // Java creates FJP threads without caller's access control context, thus we reduce
    // the surface for the leak.
    //
    // Null this (submitting) thread's contextClassLoader around execute(): if the common pool
    // spawns a fresh worker to run the cleanup task, the worker inherits the submitter's
    // contextClassLoader at birth. Leaving that classloader (which may be a web application's)
    // on a worker that lives for the JVM's lifetime would pin it and leak it
    // (https://github.com/pgjdbc/pgjdbc/issues/3953). We restore the value right after, so the
    // change is scoped to our own call: we never mutate a shared common-pool worker, whose
    // contextClassLoader unrelated tasks rely on (https://github.com/pgjdbc/pgjdbc/issues/4155).
    Thread current = Thread.currentThread();
    ClassLoader prevContextCL = current.getContextClassLoader();
    try {
      current.setContextClassLoader(null);
    } catch (SecurityException ignore) {
      // setContextClassLoader can be denied even without a SecurityManager: when this is an
      // InnocuousForkJoinWorkerThread (a registration made from a common-pool task), it throws
      // unconditionally. This is best-effort, so submit anyway with the original loader.
    }
    try {
      ForkJoinPool.commonPool().execute(
          () -> {
            RefQueueBlocker<Object> blocker =
                new RefQueueBlocker<>(queue, threadName, threadTtl);
            while (true) {
              try {
                ForkJoinPool.managedBlock(blocker);
                Node<?> ref = (Node<?>) blocker.drainOne();
                if (ref != null) {
                  ref.onClean(true);
                } else if (checkEmpty()) {
                  // The blocker waited the full threadTtl without receiving a ref, and there are
                  // no pending registrations, so the cleanup task can terminate. Keeping the task
                  // alive across transient empties amortizes ForkJoinPool submit/compensation
                  // overhead under bursty workloads.
                  break;
                }
              } catch (InterruptedException e) {
                if (checkEmpty()) {
                  LOGGER.log(Level.FINE, "Got interrupt and the cleanup queue is empty, will terminate the cleanup thread");
                  break;
                }
                LOGGER.log(Level.FINE, "Got interrupt and the cleanup queue is NOT empty. Will ignore the interrupt");
              } catch (Throwable e) {
                LOGGER.log(Level.WARNING, "Unexpected exception while executing onClean", e);
              }
            }
          }
      );
    } finally {
      try {
        current.setContextClassLoader(prevContextCL);
      } catch (SecurityException ignore) {
        // Denied by an InnocuousForkJoinWorkerThread or a SecurityManager, as above; the loader
        // was never changed, so there is nothing to restore.
      }
    }
    return true;
  }

  private synchronized boolean remove(Node<?> node) {
    // If already removed, do nothing
    if (node.next == node) {
      return false;
    }

    // Update list
    if (first == node) {
      first = node.next;
    }
    if (node.next != null) {
      node.next.prev = node.prev;
    }
    if (node.prev != null) {
      node.prev.next = node.next;
    }

    // Indicate removal by pointing the cleaner to itself
    node.next = node;
    node.prev = node;

    return true;
  }

  private class Node<T extends Throwable> extends PhantomReference<Object> implements Cleanable<T>,
      CleaningAction<T> {
    private final CleaningAction<T> action;
    private /* @Nullable */ Node<?> prev;
    private /* @Nullable */ Node<?> next;

    Node(Object referent, CleaningAction<T> action) {
      super(referent, queue);
      this.action = action;
      //Objects.requireNonNull(referent); // poor man`s reachabilityFence
    }

    @Override
    public void clean() throws T {
      onClean(false);
    }

    @Override
    public void onClean(boolean leak) throws T {
      if (!remove(this)) {
        return;
      }
      action.onClean(leak);
    }
  }
}
