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


package org.apache.spark.util

import java.util.concurrent.{CountDownLatch, TimeUnit}

import scala.concurrent.Future
import scala.concurrent.duration._
import scala.util.Random

import org.scalatest.concurrent.Eventually._

import org.apache.spark.{SparkException, SparkFunSuite, SparkThrowable}

class ThreadUtilsSuite extends SparkFunSuite {

  test("newDaemonSingleThreadExecutor") {
    val executor = ThreadUtils.newDaemonSingleThreadExecutor("this-is-a-thread-name")
    @volatile var threadName = ""
    executor.submit(new Runnable {
      override def run(): Unit = {
        threadName = Thread.currentThread().getName()
      }
    })
    executor.shutdown()
    executor.awaitTermination(10, TimeUnit.SECONDS)
    assert(threadName === "this-is-a-thread-name")
  }

  test("newDaemonSingleThreadScheduledExecutor") {
    val executor = ThreadUtils.newDaemonSingleThreadScheduledExecutor("this-is-a-thread-name")
    try {
      val latch = new CountDownLatch(1)
      @volatile var threadName = ""
      executor.schedule(new Runnable {
        override def run(): Unit = {
          threadName = Thread.currentThread().getName()
          latch.countDown()
        }
      }, 1, TimeUnit.MILLISECONDS)
      latch.await(10, TimeUnit.SECONDS)
      assert(threadName === "this-is-a-thread-name")
    } finally {
      executor.shutdownNow()
    }
  }

  test("newDaemonCachedThreadPool") {
    val maxThreadNumber = 10
    val startThreadsLatch = new CountDownLatch(maxThreadNumber)
    val latch = new CountDownLatch(1)
    val cachedThreadPool = ThreadUtils.newDaemonCachedThreadPool(
      "ThreadUtilsSuite-newDaemonCachedThreadPool",
      maxThreadNumber,
      keepAliveSeconds = 2)
    try {
      for (_ <- 1 to maxThreadNumber) {
        cachedThreadPool.execute(() => {
          startThreadsLatch.countDown()
          latch.await(10, TimeUnit.SECONDS)
        })
      }
      startThreadsLatch.await(10, TimeUnit.SECONDS)
      assert(cachedThreadPool.getActiveCount === maxThreadNumber)
      assert(cachedThreadPool.getQueue.size === 0)

      // Submit a new task and it should be put into the queue since the thread number reaches the
      // limitation
      cachedThreadPool.execute(() => latch.await(10, TimeUnit.SECONDS))

      assert(cachedThreadPool.getActiveCount === maxThreadNumber)
      assert(cachedThreadPool.getQueue.size === 1)

      latch.countDown()
      eventually(timeout(10.seconds)) {
        // All threads should be stopped after keepAliveSeconds
        assert(cachedThreadPool.getActiveCount === 0)
        assert(cachedThreadPool.getPoolSize === 0)
      }
    } finally {
      cachedThreadPool.shutdownNow()
    }
  }

  test("newDaemonBlockingThreadPoolExecutorService") {
    val nThread = 3
    val workQueueSize = 5
    val submithreadsLatch = new CountDownLatch(nThread + workQueueSize + 1)
    val latch = new CountDownLatch(1)
    val blockingPool = ThreadUtils.newDaemonBlockingThreadPoolExecutorService(
      nThread, workQueueSize, "ThreadUtilsSuite-newDaemonBlockingThreadPoolExecutorService")

    try {
      val submitThread = new Thread(() => {
        (0 until nThread + workQueueSize + 1).foreach { i =>
          blockingPool.execute(() => {
            latch.await(10, TimeUnit.SECONDS)
          })
          submithreadsLatch.countDown()
        }
      })
      submitThread.setDaemon(true)
      submitThread.start()

      // the last one task submission will be blocked until previous tasks completed
      eventually(timeout(10.seconds)) {
        assert(submithreadsLatch.getCount === 1L)
      }
      latch.countDown()
      eventually(timeout(10.seconds)) {
        assert(submithreadsLatch.getCount === 0L)
        assert(!submitThread.isAlive)
      }
    } finally {
      blockingPool.shutdownNow()
    }
  }

  test("sameThread") {
    val callerThreadName = Thread.currentThread().getName()
    val f = Future {
      Thread.currentThread().getName()
    }(ThreadUtils.sameThread)
    val futureThreadName = ThreadUtils.awaitResult(f, 10.seconds)
    assert(futureThreadName === callerThreadName)
  }

  test("runInNewThread") {
    import ThreadUtils._
    assert(runInNewThread("thread-name") { Thread.currentThread().getName } === "thread-name")
    assert(runInNewThread("thread-name") { Thread.currentThread().isDaemon })
    assert(
      runInNewThread("thread-name", isDaemon = false) { Thread.currentThread().isDaemon } === false
    )
    val uniqueExceptionMessage = "test" + Random.nextInt()
    val exception = intercept[IllegalArgumentException] {
      runInNewThread("thread-name") { throw new IllegalArgumentException(uniqueExceptionMessage) }
    }
    assert(exception.getMessage === uniqueExceptionMessage)
    val stacktrace = exception.getStackTrace.mkString("\n")
    assert(stacktrace.contains(
      "... run in separate thread using org.apache.spark.util.ThreadUtils ..."),
      "stack trace does not contain expected place holder"
    )
    assert(!stacktrace.contains("ThreadUtils.scala"),
      "stack trace contains unexpected references to ThreadUtils"
    )
  }

  test("SPARK-47833: wrapCallerStacktrace") {
    var runnerThreadName: String = null
    var exception: Throwable = null
    val t = new Thread() {
      override def run(): Unit = {
        runnerThreadName = Thread.currentThread().getName
        internalMethod()
      }
      private def internalMethod(): Unit = {
        throw new RuntimeException(s"Error occurred on $runnerThreadName")
      }
    }
    t.setDaemon(true)
    t.setUncaughtExceptionHandler { case (_, e) => exception = e }
    t.start()
    t.join()

    ThreadUtils.wrapCallerStacktrace(exception, s"run in separate thread: $runnerThreadName")

    val stacktrace = exception.getStackTrace.mkString("\n")
    assert(stacktrace.contains("internalMethod"),
      "stack trace does not contain real exception stack trace"
    )
    assert(stacktrace.contains(s"... run in separate thread: $runnerThreadName ..."),
      "stack trace does not contain expected place holder"
    )
    assert(stacktrace.contains("org.scalatest.Suite.run"),
      "stack trace does not contain caller stack trace"
    )
    assert(!stacktrace.contains("ThreadUtils.scala"),
      "stack trace contains unexpected references to ThreadUtils"
    )
  }

  test("parmap should be interruptible") {
    val t = new Thread() {
      setDaemon(true)

      override def run(): Unit = {
        try {
          // "par" is uninterruptible. The following will keep running even if the thread is
          // interrupted. We should prefer to use "ThreadUtils.parmap".
          //
          // (1 to 10).par.flatMap { i =>
          //   Thread.sleep(100000)
          //   1 to i
          // }
          //
          ThreadUtils.parmap(1 to 10, "test", 2) { i =>
            Thread.sleep(100000)
            1 to i
          }.flatten
        } catch {
          case _: InterruptedException => // excepted
        }
      }
    }
    t.start()
    eventually(timeout(10.seconds)) {
      assert(t.isAlive)
    }
    t.interrupt()
    eventually(timeout(10.seconds)) {
      assert(!t.isAlive)
    }
  }

  test("awaitResult preserves SparkThrowable when flag is true") {
    import java.io.IOException

    val sparkThrowableEx = new RuntimeException("structured error") with SparkThrowable {
      override def getCondition: String = "TEST_ERROR_CLASS"
      override def getMessageParameters: java.util.Map[String, String] =
        java.util.Collections.emptyMap()
    }

    // With preserveSparkThrowable=true, SparkThrowable is re-thrown directly.
    val f1 = Future {
      throw sparkThrowableEx
    }(ThreadUtils.sameThread)
    val caught1 = intercept[RuntimeException] {
      ThreadUtils.awaitResult(f1, 1.seconds, preserveSparkThrowable = true)
    }
    assert(caught1.isInstanceOf[SparkThrowable])
    assert(caught1.asInstanceOf[SparkThrowable].getCondition == "TEST_ERROR_CLASS")
    assert(caught1.getSuppressed.nonEmpty)

    // With preserveSparkThrowable=false (default), SparkThrowable is wrapped in SparkException.
    val f2 = Future {
      throw sparkThrowableEx
    }(ThreadUtils.sameThread)
    val caught2 = intercept[SparkException] {
      ThreadUtils.awaitResult(f2, 1.seconds)
    }
    assert(caught2.getCause.isInstanceOf[SparkThrowable])

    // Plain exceptions are always wrapped regardless of the flag.
    val plainEx = new IOException("plain error")
    val f3 = Future {
      throw plainEx
    }(ThreadUtils.sameThread)
    val caught3 = intercept[SparkException] {
      ThreadUtils.awaitResult(f3, 1.seconds, preserveSparkThrowable = true)
    }
    assert(caught3.getCause eq plainEx)
  }

  test("awaitResult (JFuture) preserves SparkThrowable when flag is true") {
    val sparkThrowableEx = new RuntimeException("structured error") with SparkThrowable {
      override def getCondition: String = "TEST_ERROR_CLASS"
      override def getMessageParameters: java.util.Map[String, String] =
        java.util.Collections.emptyMap()
    }

    // scalastyle:off sparkThreadPools
    val jfuture = new java.util.concurrent.CompletableFuture[String]()
    // scalastyle:on sparkThreadPools
    jfuture.completeExceptionally(sparkThrowableEx)

    val caught = intercept[RuntimeException] {
      ThreadUtils.awaitResult(jfuture, 10.seconds, preserveSparkThrowable = true)
    }
    assert(caught.isInstanceOf[SparkThrowable])
    assert(caught.asInstanceOf[SparkThrowable].getCondition == "TEST_ERROR_CLASS")
    assert(caught.getSuppressed.nonEmpty)
  }

  test("parmap preserves SparkThrowable when flag is true") {
    val sparkThrowableEx = new RuntimeException("structured error") with SparkThrowable {
      override def getCondition: String = "TEST_ERROR_CLASS"
      override def getMessageParameters: java.util.Map[String, String] =
        java.util.Collections.emptyMap()
    }

    // With preserveSparkThrowable=true, the original SparkThrowable is re-thrown.
    val caught1 = intercept[RuntimeException] {
      ThreadUtils.parmap(Seq(1), "test", 1, preserveSparkThrowable = true) { _ =>
        throw sparkThrowableEx
      }
    }
    assert(caught1.isInstanceOf[SparkThrowable])
    assert(caught1.asInstanceOf[SparkThrowable].getCondition == "TEST_ERROR_CLASS")
    assert(caught1.getSuppressed.nonEmpty)

    // With preserveSparkThrowable=false, it is wrapped in SparkException.
    val caught2 = intercept[SparkException] {
      ThreadUtils.parmap(Seq(1), "test", 1, preserveSparkThrowable = false) { _ =>
        throw sparkThrowableEx
      }
    }
    assert(caught2.getCause.isInstanceOf[SparkThrowable])
  }
}
