<commit_before>/*
 * Copyright 2009 LinkedIn, Inc
 *
 * 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 com.linkedin.norbert.network

import com.linkedin.norbert.cluster._
import com.linkedin.norbert.util.Logging
import com.google.protobuf.Message

/**
 * A component which provides a network client for interacting with nodes in a cluster.
 */
trait NetworkClientFactoryComponent {
  this: ChannelPoolComponent with ClusterComponent with RouterFactoryComponent =>

  val networkClientFactory: NetworkClientFactory

  /**
   * The network client interface for interacting with nodes in a cluster.
   */
  trait NetworkClient {
    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     *
     * @return a <code>ResponseIterator</code>. One response will be returned by each <code>Node</code>
     * the message was sent to.
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage(ids: Seq[Id], message: Message): ResponseIterator
    
    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param messageCustomizer a callback method which allows the user to customize the <code>Message</code>
     * before it is sent to the <code>Node</code>. The callback will receive the original message passed to <code>sendMessage</code>
     * the <code>Node</code> the request is being sent to and the <code>Id</code>s which reside on that
     * <code>Node</code>. The callback should return a <code>Message</code> which has been customized.
     *
     * @return a <code>ResponseIterator</code>. One response will be returned by each <code>Node</code>
     * the message was sent to.
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage(ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message): ResponseIterator

    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param responseAggregator a callback method which allows the user to aggregate all the responses
     * and return a single object to the caller.  The callback will receive the original message passed to
     * <code>sendMessage</code> and the <code>ResponseIterator</code> for the request.
     *
     * @return the return value of the <code>responseAggregator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage[A](ids: Seq[Id], message: Message, responseAggregator: (Message, ResponseIterator) => A): A

    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param messageCustomizer a callback method which allows the user to customize the <code>Message</code>
     * before it is sent to the <code>Node</code>. The callback will receive the original message passed to <code>sendMessage</code>
     * the <code>Node</code> the request is being sent to and the <code>Id</code>s which reside on that
     * <code>Node</code>. The callback should return a <code>Message</code> which has been customized.
     * @param responseAggregator a callback method which allows the user to aggregate all the responses
     * and return a single object to the caller.  The callback will receive the original message passed to
     * <code>sendMessage</code> and the <code>ResponseIterator</code> for the request.
     *
     * @return the return value of the <code>responseAggregator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage[A](ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message,
                responseAggregator: (Message, ResponseIterator) => A): A

    /**
     * Sends a <code>Message</code> to the specified <code>Node</code>.
     *
     * @param node the <code>Node</code> to send the message
     * @param message the <code>Message</code> to send
     *
     * @return a <code>ResponseIterator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     * @throws InvalidNodeException thrown if an attempt is made to send a message to an unavailable <code>Node</code>
     */
    def sendMessageToNode(node: Node, message: Message): ResponseIterator

    /**
     * Queries whether or not a connection to the cluster is established.
     *
     * @return true if connected, false otherwise
     */
    def isConnected: Boolean

    /**
     * Closes the <code>NetworkClient</code> and releases resources held.
     * 
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def close: Unit
  }

  /**
   * Factory for creating instances of <code>NetworkClient</code>.
   */
  class NetworkClientFactory extends Logging {

    /**
     * Create a new <code>NetworkClient</code>.
     *
     * @return a new <code>NetworkClient</code>
     */
    def newNetworkClient: NetworkClient = {
      cluster.awaitConnectionUninterruptibly

      val client = new NetworkClient with ClusterListener {
        @volatile private var routerOption: Option[Router] = cluster.router
        @volatile private var shutdownSwitch = false
    
        def sendMessage(ids: Seq[Id], message: Message): ResponseIterator = doIfNotShutdown {
          val nodes = calculateNodesFromIds(ids)
          log.ifDebug("Sending message [%s] to nodes: %s", message, nodes)
          val request = Request(message, nodes.size)
          channelPool.sendRequest(nodes.keySet, request)
          request.responseIterator
        }

        def sendMessage(ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message): ResponseIterator = doIfNotShutdown {
          val nodes = calculateNodesFromIds(ids)
          val responseIterator = new Request.ResponseIteratorImpl(nodes.size)

          nodes.foreach { case (node, r) =>
            val m = messageCustomizer(message, node, r)
            log.ifDebug("Sending message [%s] to node: %s", message, node)
            val request = Request(message, nodes.size, responseIterator)<commit_msg>Use the correctly customized message when sending request<commit_after>/*
 * Copyright 2009 LinkedIn, Inc
 *
 * 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 com.linkedin.norbert.network

import com.linkedin.norbert.cluster._
import com.linkedin.norbert.util.Logging
import com.google.protobuf.Message

/**
 * A component which provides a network client for interacting with nodes in a cluster.
 */
trait NetworkClientFactoryComponent {
  this: ChannelPoolComponent with ClusterComponent with RouterFactoryComponent =>

  val networkClientFactory: NetworkClientFactory

  /**
   * The network client interface for interacting with nodes in a cluster.
   */
  trait NetworkClient {
    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     *
     * @return a <code>ResponseIterator</code>. One response will be returned by each <code>Node</code>
     * the message was sent to.
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage(ids: Seq[Id], message: Message): ResponseIterator
    
    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param messageCustomizer a callback method which allows the user to customize the <code>Message</code>
     * before it is sent to the <code>Node</code>. The callback will receive the original message passed to <code>sendMessage</code>
     * the <code>Node</code> the request is being sent to and the <code>Id</code>s which reside on that
     * <code>Node</code>. The callback should return a <code>Message</code> which has been customized.
     *
     * @return a <code>ResponseIterator</code>. One response will be returned by each <code>Node</code>
     * the message was sent to.
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage(ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message): ResponseIterator

    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param responseAggregator a callback method which allows the user to aggregate all the responses
     * and return a single object to the caller.  The callback will receive the original message passed to
     * <code>sendMessage</code> and the <code>ResponseIterator</code> for the request.
     *
     * @return the return value of the <code>responseAggregator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage[A](ids: Seq[Id], message: Message, responseAggregator: (Message, ResponseIterator) => A): A

    /**
     * Sends a <code>Message</code> to the specified <code>Id</code>s. The <code>NetworkClient</code>
     * will interact with the <code>Cluster</code> to calculate which <code>Node</code>s the message
     * must be sent to.
     *
     * @param ids the <code>Id</code>s to which the message is addressed
     * @param message the message to send
     * @param messageCustomizer a callback method which allows the user to customize the <code>Message</code>
     * before it is sent to the <code>Node</code>. The callback will receive the original message passed to <code>sendMessage</code>
     * the <code>Node</code> the request is being sent to and the <code>Id</code>s which reside on that
     * <code>Node</code>. The callback should return a <code>Message</code> which has been customized.
     * @param responseAggregator a callback method which allows the user to aggregate all the responses
     * and return a single object to the caller.  The callback will receive the original message passed to
     * <code>sendMessage</code> and the <code>ResponseIterator</code> for the request.
     *
     * @return the return value of the <code>responseAggregator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def sendMessage[A](ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message,
                responseAggregator: (Message, ResponseIterator) => A): A

    /**
     * Sends a <code>Message</code> to the specified <code>Node</code>.
     *
     * @param node the <code>Node</code> to send the message
     * @param message the <code>Message</code> to send
     *
     * @return a <code>ResponseIterator</code>
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     * @throws InvalidNodeException thrown if an attempt is made to send a message to an unavailable <code>Node</code>
     */
    def sendMessageToNode(node: Node, message: Message): ResponseIterator

    /**
     * Queries whether or not a connection to the cluster is established.
     *
     * @return true if connected, false otherwise
     */
    def isConnected: Boolean

    /**
     * Closes the <code>NetworkClient</code> and releases resources held.
     * 
     * @throws ClusterShutdownException thrown if the cluster is shutdown when the method is called
     */
    def close: Unit
  }

  /**
   * Factory for creating instances of <code>NetworkClient</code>.
   */
  class NetworkClientFactory extends Logging {

    /**
     * Create a new <code>NetworkClient</code>.
     *
     * @return a new <code>NetworkClient</code>
     */
    def newNetworkClient: NetworkClient = {
      cluster.awaitConnectionUninterruptibly

      val client = new NetworkClient with ClusterListener {
        @volatile private var routerOption: Option[Router] = cluster.router
        @volatile private var shutdownSwitch = false
    
        def sendMessage(ids: Seq[Id], message: Message): ResponseIterator = doIfNotShutdown {
          val nodes = calculateNodesFromIds(ids)
          log.ifDebug("Sending message [%s] to nodes: %s", message, nodes)
          val request = Request(message, nodes.size)
          channelPool.sendRequest(nodes.keySet, request)
          request.responseIterator
        }

        def sendMessage(ids: Seq[Id], message: Message, messageCustomizer: (Message, Node, Seq[Id]) => Message): ResponseIterator = doIfNotShutdown {
          val nodes = calculateNodesFromIds(ids)
          val responseIterator = new Request.ResponseIteratorImpl(nodes.size)

          nodes.foreach { case (node, r) =>
            val m = messageCustomizer(message, node, r)
            log.ifDebug("Sending message [%s] to node: %s", m, node)
            val request = Request(m, nodes.size, responseIterator)
{'fraction_non_alphanumeric': 0.09283076531261938, 'fraction_numerical': 0.0010187189609066599, 'mean_word_length': 3.3269972451790633, 'pattern_counts': {'":': 0, '<': 235, '<?xml version=': 0, '>': 249, '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': '2281463', 'n_tokens_mistral': 4191, 'n_tokens_neox': 4114, 'n_words': 1654}