<commit_before>package com.softwaremill.mqperf.mq

import java.util.concurrent.TimeUnit

import com.mongodb._
import com.mongodb.client.model.{CreateCollectionOptions, UpdateOptions}
import com.mongodb.client.{MongoCollection, MongoCursor, MongoDatabase}
import org.bson.Document
import org.bson.types.ObjectId

import scala.annotation.tailrec
import scala.collection.JavaConverters._
import scala.concurrent.Future

class MongoCappedMq(configMap: Map[String, String]) extends Mq {

  import MongoCappedMq._

  override type MsgId = ObjectId

  private val client = new MongoClient(configMap("host"))

  private val concern = if (configMap("write_concern") == "replica")
    WriteConcern.W2 else WriteConcern.ACKNOWLEDGED

  private val db: MongoDatabase = client.getDatabase("mq").withWriteConcern(concern)

  private val queueColl: MongoCollection[Document] = {
    if (!db.listCollectionNames().iterator().asScala.contains(QueueCollectionName)) {
      val options = new CreateCollectionOptions()
        .capped(true)
        .sizeInBytes(configMap("queue_size_in_bytes").toLong)
      db.createCollection(QueueCollectionName, options)
    }
    db.getCollection(QueueCollectionName)
  }

  private val cursorStateColl: MongoCollection[Document] = {
    if (!db.listCollectionNames().iterator().asScala.contains(CursorStateCollectionName)) {
      db.createCollection(CursorStateCollectionName)
    }
    db.getCollection(CursorStateCollectionName)
  }

  override def createSender(): MqSender = new MongoCappedMqSender(queueColl)

  override def createReceiver(): MqReceiver = new MongoCappedMqReceiver(queueColl, cursorStateColl)

  class MongoCappedMqSender(queueColl: MongoCollection[Document]) extends MqSender {

    /**
     * Synchronous - must wait for the messages to be sent
     */
    override def send(msgs: List[String]): Unit = {
      val messagesToInsert: java.util.List[Document] = msgs.map { msg =>
        new Document(MessageField, msg)
      }.asJava
      queueColl.insertMany(messagesToInsert)
    }
  }

  class MongoCappedMqReceiver(queueColl: MongoCollection[Document], cursorStateColl: MongoCollection[Document]) extends MqReceiver {

    private val nodeId: Int = ObjectId.getGeneratedMachineIdentifier

    private val cursor: MongoCursor[Document] = {
      val filter = Option(cursorStateColl.find(new BasicDBObject(NodeIdField, nodeId)).first())
        .map(_.getObjectId(LastDocField))
        .map(lastDocId => new BasicDBObject(IdField, new BasicDBObject("$gt", lastDocId)))
        .getOrElse(new BasicDBObject())

      // TODO handle lost capped position
      queueColl
        .find(filter)
        .sort(new BasicDBObject("$natural", 1))
        .cursorType(CursorType.TailableAwait)
        .maxAwaitTime(configMap("await_ms").toInt, TimeUnit.MILLISECONDS)
        .batchSize(configMap("batch_size").toInt)
        .iterator()
    }

    override def receive(maxMsgCount: Int): List[(ObjectId, String)] = {
      doReceive(Nil, maxMsgCount)
    }

    @tailrec
    private def doReceive(acc: List[(ObjectId, String)], limit: Int): List[(ObjectId, String)] = {
      if (limit > 0) {
        Option(cursor.tryNext()) match {
          case Some(doc) =>
            val newElement = doc.getObjectId(IdField) -> doc.getString(MessageField)
            doReceive(newElement :: acc, limit - 1)
          case None =>
            acc
        }
      }
      else {
        acc
      }
    }

    /**
     * Can be asynchronous
     */
    override def ack(ids: List[ObjectId]): Unit = {
      import scala.concurrent.ExecutionContext.Implicits.global
      Future {
        val lastSeenId = ids.max
        cursorStateColl.updateOne(new BasicDBObject(NodeIdField, nodeId), new BasicDBObject("$set", new BasicDBObject(LastDocField, lastSeenId)), new UpdateOptions().upsert(true))
      }<commit_msg>Avoid Option allocation in MongoCappedMqReceiver<commit_after>package com.softwaremill.mqperf.mq

import java.util.concurrent.TimeUnit

import com.mongodb._
import com.mongodb.client.model.{CreateCollectionOptions, UpdateOptions}
import com.mongodb.client.{MongoCollection, MongoCursor, MongoDatabase}
import org.bson.Document
import org.bson.types.ObjectId

import scala.annotation.tailrec
import scala.collection.JavaConverters._
import scala.concurrent.Future

class MongoCappedMq(configMap: Map[String, String]) extends Mq {

  import MongoCappedMq._

  override type MsgId = ObjectId

  private val client = new MongoClient(configMap("host"))

  private val concern = if (configMap("write_concern") == "replica")
    WriteConcern.W2 else WriteConcern.ACKNOWLEDGED

  private val db: MongoDatabase = client.getDatabase("mq").withWriteConcern(concern)

  private val queueColl: MongoCollection[Document] = {
    if (!db.listCollectionNames().iterator().asScala.contains(QueueCollectionName)) {
      val options = new CreateCollectionOptions()
        .capped(true)
        .sizeInBytes(configMap("queue_size_in_bytes").toLong)
      db.createCollection(QueueCollectionName, options)
    }
    db.getCollection(QueueCollectionName)
  }

  private val cursorStateColl: MongoCollection[Document] = {
    if (!db.listCollectionNames().iterator().asScala.contains(CursorStateCollectionName)) {
      db.createCollection(CursorStateCollectionName)
    }
    db.getCollection(CursorStateCollectionName)
  }

  override def createSender(): MqSender = new MongoCappedMqSender(queueColl)

  override def createReceiver(): MqReceiver = new MongoCappedMqReceiver(queueColl, cursorStateColl)

  class MongoCappedMqSender(queueColl: MongoCollection[Document]) extends MqSender {

    /**
     * Synchronous - must wait for the messages to be sent
     */
    override def send(msgs: List[String]): Unit = {
      val messagesToInsert: java.util.List[Document] = msgs.map { msg =>
        new Document(MessageField, msg)
      }.asJava
      queueColl.insertMany(messagesToInsert)
    }
  }

  class MongoCappedMqReceiver(queueColl: MongoCollection[Document], cursorStateColl: MongoCollection[Document]) extends MqReceiver {

    private val nodeId: Int = ObjectId.getGeneratedMachineIdentifier

    private val cursor: MongoCursor[Document] = {
      val existingCursorState = cursorStateColl.find(new BasicDBObject(NodeIdField, nodeId)).first()
      val filter = if (existingCursorState == null) {
        new BasicDBObject()
      }
      else {
        val lastDocId = existingCursorState.getObjectId(LastDocField)
        new BasicDBObject(IdField, new BasicDBObject("$gt", lastDocId))
      }

      // TODO handle lost capped position
      queueColl
        .find(filter)
        .sort(new BasicDBObject("$natural", 1))
        .cursorType(CursorType.TailableAwait)
        .maxAwaitTime(configMap("await_ms").toInt, TimeUnit.MILLISECONDS)
        .batchSize(configMap("batch_size").toInt)
        .iterator()
    }

    override def receive(maxMsgCount: Int): List[(ObjectId, String)] = {
      doReceive(Nil, maxMsgCount)
    }

    @tailrec
    private def doReceive(acc: List[(ObjectId, String)], limit: Int): List[(ObjectId, String)] = {
      if (limit > 0) {
        val msg = cursor.tryNext()
        if (msg == null) {
          acc
        }
        else {
          val newElement = msg.getObjectId(IdField) -> msg.getString(MessageField)
          doReceive(newElement :: acc, limit - 1)
        }
      }
      else {
        acc
      }
    }

    /**
     * Can be asynchronous
     */
    override def ack(ids: List[ObjectId]): Unit = {
      import scala.concurrent.ExecutionContext.Implicits.global
      Future {
        val lastSeenId = ids.max
        cursorStateColl.updateOne(new BasicDBObject(NodeIdField, nodeId), new BasicDBObject("$set", new BasicDBObject(LastDocField, lastSeenId)), new UpdateOptions().upsert(true))
      }
{'fraction_non_alphanumeric': 0.10422282120395328, 'fraction_numerical': 0.0010268258246694903, 'mean_word_length': 3.6215895610913407, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '5999016', 'n_tokens_mistral': 2344, 'n_tokens_neox': 2236, 'n_words': 464}