<commit_before>/*
 * 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.flink.streaming.api.scala

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

import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.SinkFunction
import org.apache.flink.streaming.api.scala.AsyncDataStreamITCase._
import org.apache.flink.streaming.api.scala.async.{ResultFuture, RichAsyncFunction}
import org.apache.flink.test.util.AbstractTestBase
import org.junit.Assert._
import org.junit.Test

import scala.collection.mutable
import scala.concurrent.{ExecutionContext, Future}

object AsyncDataStreamITCase {
  private var testResult: mutable.ArrayBuffer[Int] = _
}

class AsyncDataStreamITCase extends AbstractTestBase {

  @Test
  def testOrderedWait(): Unit = {
    testAsyncWait(true)
  }

  @Test
  def testUnorderedWait(): Unit = {
    testAsyncWait(false)
  }

  private def testAsyncWait(ordered: Boolean): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1, 2)

    val timeout = 1L
    val asyncMapped = if (ordered) {
      AsyncDataStream.orderedWait(
        source, new AsyncFunctionWithTimeoutExpired(), timeout, TimeUnit.MILLISECONDS)
    } else {
      AsyncDataStream.unorderedWait(
        source, new AsyncFunctionWithTimeoutExpired(), timeout, TimeUnit.MILLISECONDS)
    }

    executeAndValidate(ordered, env, asyncMapped, mutable.ArrayBuffer[Int](2, 6))
  }

  private def executeAndValidate(ordered: Boolean,
      env: StreamExecutionEnvironment,
      dataStream: DataStream[Int],
      expectedResult: mutable.ArrayBuffer[Int]): Unit = {

    testResult = mutable.ArrayBuffer[Int]()
    dataStream.addSink(new SinkFunction[Int]() {
      override def invoke(value: Int) {
        testResult += value
      }
    })

    env.execute("testAsyncDataStream")

    if (ordered) {
      assertEquals(expectedResult, testResult)
    } else {
      assertEquals(expectedResult, testResult.sorted)
    }
  }

  @Test
  def testOrderedWaitUsingAnonymousFunction(): Unit = {
    testAsyncWaitUsingAnonymousFunction(true)
  }

  @Test
  def testUnorderedWaitUsingAnonymousFunction(): Unit = {
    testAsyncWaitUsingAnonymousFunction(false)
  }

  @Test
  def testRichAsyncFunctionRuntimeContext(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1)

    val timeout = 10000L
    val richAsyncFunction = new MyRichAsyncFunction
    val asyncMapped = AsyncDataStream
      .unorderedWait(source, richAsyncFunction, timeout, TimeUnit.MILLISECONDS)

    executeAndValidate(false, env, asyncMapped, mutable.ArrayBuffer[Int](2))
  }

  private def testAsyncWaitUsingAnonymousFunction(ordered: Boolean): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1, 2)

    val asyncFunction: (Int, ResultFuture[Int]) => Unit =
      (input, collector: ResultFuture[Int]) => Future {
          collector.complete(Seq(input * 2))
      }(ExecutionContext.global)

    val timeout = 10000L
    val asyncMapped = if (ordered) {
      AsyncDataStream.orderedWait(source, timeout, TimeUnit.MILLISECONDS) {
        asyncFunction
      }
    } else {
      AsyncDataStream.unorderedWait(source, timeout, TimeUnit.MILLISECONDS) {
        asyncFunction
      }
    }

    executeAndValidate(ordered, env, asyncMapped, mutable.ArrayBuffer[Int](2, 4))
  }

}

class AsyncFunctionWithTimeoutExpired extends RichAsyncFunction[Int, Int] {
  @transient var timeoutLatch: CountDownLatch = _
  @transient var invokeLatch: CountDownLatch = _

  override def open(parameters: Configuration): Unit = {
    timeoutLatch = new CountDownLatch(1)
    invokeLatch = new CountDownLatch(1)
  }

  override def asyncInvoke(input: Int, resultFuture: ResultFuture[Int]): Unit = {
    Future {
      // trigger the timeout of the even input number
      if (input % 2 == 0) {
        invokeLatch.await()
        resultFuture.complete(Seq(input * 2))
      } else {
        resultFuture.complete(Seq(input * 2))
        timeoutLatch.countDown()
      }
    } (ExecutionContext.global)
  }
  override def timeout(input: Int, resultFuture: ResultFuture[Int]): Unit = {
    if (input % 2 == 0) {
      resultFuture.complete(Seq(input * 3))
      invokeLatch.countDown()
    } else {
      timeoutLatch.await()
      resultFuture.complete(Seq(input * 3))
    }
  }
}

class MyRichAsyncFunction extends RichAsyncFunction[Int, Int] {

  override def open(parameters: Configuration): Unit = {
    assertEquals(getRuntimeContext.getNumberOfParallelSubtasks, 1)<commit_msg>[FLINK-13605][tests] Fix unstable case AsyncDataStreamITCase#testUnorderedWait<commit_after>/*
 * 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.flink.streaming.api.scala

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

import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.SinkFunction
import org.apache.flink.streaming.api.scala.AsyncDataStreamITCase._
import org.apache.flink.streaming.api.scala.async.{ResultFuture, RichAsyncFunction}
import org.apache.flink.test.util.AbstractTestBase
import org.junit.Assert._
import org.junit.Test

import scala.collection.mutable
import scala.concurrent.{ExecutionContext, Future}

object AsyncDataStreamITCase {
  private var testResult: mutable.ArrayBuffer[Int] = _
}

class AsyncDataStreamITCase extends AbstractTestBase {

  @Test
  def testOrderedWait(): Unit = {
    testAsyncWait(true)
  }

  @Test
  def testUnorderedWait(): Unit = {
    testAsyncWait(false)
  }

  private def testAsyncWait(ordered: Boolean): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1)

    val timeout = 1L
    val asyncMapped = if (ordered) {
      AsyncDataStream.orderedWait(
        source, new AsyncFunctionWithTimeoutExpired(), timeout, TimeUnit.MILLISECONDS)
    } else {
      AsyncDataStream.unorderedWait(
        source, new AsyncFunctionWithTimeoutExpired(), timeout, TimeUnit.MILLISECONDS)
    }

    executeAndValidate(ordered, env, asyncMapped, mutable.ArrayBuffer[Int](3))
  }

  private def executeAndValidate(ordered: Boolean,
      env: StreamExecutionEnvironment,
      dataStream: DataStream[Int],
      expectedResult: mutable.ArrayBuffer[Int]): Unit = {

    testResult = mutable.ArrayBuffer[Int]()
    dataStream.addSink(new SinkFunction[Int]() {
      override def invoke(value: Int) {
        testResult += value
      }
    })

    env.execute("testAsyncDataStream")

    if (ordered) {
      assertEquals(expectedResult, testResult)
    } else {
      assertEquals(expectedResult, testResult.sorted)
    }
  }

  @Test
  def testOrderedWaitUsingAnonymousFunction(): Unit = {
    testAsyncWaitUsingAnonymousFunction(true)
  }

  @Test
  def testUnorderedWaitUsingAnonymousFunction(): Unit = {
    testAsyncWaitUsingAnonymousFunction(false)
  }

  @Test
  def testRichAsyncFunctionRuntimeContext(): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1)

    val timeout = 10000L
    val richAsyncFunction = new MyRichAsyncFunction
    val asyncMapped = AsyncDataStream
      .unorderedWait(source, richAsyncFunction, timeout, TimeUnit.MILLISECONDS)

    executeAndValidate(false, env, asyncMapped, mutable.ArrayBuffer[Int](2))
  }

  private def testAsyncWaitUsingAnonymousFunction(ordered: Boolean): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    val source = env.fromElements(1, 2)

    val asyncFunction: (Int, ResultFuture[Int]) => Unit =
      (input, collector: ResultFuture[Int]) => Future {
          collector.complete(Seq(input * 2))
      }(ExecutionContext.global)

    val timeout = 10000L
    val asyncMapped = if (ordered) {
      AsyncDataStream.orderedWait(source, timeout, TimeUnit.MILLISECONDS) {
        asyncFunction
      }
    } else {
      AsyncDataStream.unorderedWait(source, timeout, TimeUnit.MILLISECONDS) {
        asyncFunction
      }
    }

    executeAndValidate(ordered, env, asyncMapped, mutable.ArrayBuffer[Int](2, 4))
  }

}

class AsyncFunctionWithTimeoutExpired extends RichAsyncFunction[Int, Int] {
  @transient var invokeLatch: CountDownLatch = _

  override def open(parameters: Configuration): Unit = {
    invokeLatch = new CountDownLatch(1)
  }

  override def asyncInvoke(input: Int, resultFuture: ResultFuture[Int]): Unit = {
    Future {
      invokeLatch.await()
      resultFuture.complete(Seq(input * 2))
    } (ExecutionContext.global)
  }
  override def timeout(input: Int, resultFuture: ResultFuture[Int]): Unit = {
    resultFuture.complete(Seq(input * 3))
    invokeLatch.countDown()
  }
}

class MyRichAsyncFunction extends RichAsyncFunction[Int, Int] {

  override def open(parameters: Configuration): Unit = {
    assertEquals(getRuntimeContext.getNumberOfParallelSubtasks, 1)
{'fraction_non_alphanumeric': 0.09024252679075014, 'fraction_numerical': 0.007144200037601053, 'mean_word_length': 3.95990675990676, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '6351525', 'n_tokens_mistral': 3102, 'n_tokens_neox': 3008, 'n_words': 754}