<commit_before>package org.broadinstitute.dsde.rawls.dataaccess

import java.util.concurrent.atomic.{AtomicInteger, AtomicBoolean}
import java.util.concurrent.locks.{Lock, StampedLock, ReadWriteLock, ReentrantReadWriteLock}

import com.tinkerpop.blueprints.Element
import com.tinkerpop.blueprints.impls.orient.OrientElement
import com.orientechnologies.orient.core.exception.OConcurrentModificationException
import com.tinkerpop.blueprints.Graph
import com.tinkerpop.blueprints.impls.orient.OrientConfigurableGraph.THREAD_MODE
import com.tinkerpop.blueprints.impls.orient.{OrientGraph, OrientGraphFactory}
import com.typesafe.scalalogging.slf4j.LazyLogging
import org.broadinstitute.dsde.rawls.RawlsException
import org.broadinstitute.dsde.rawls.model.WorkspaceName

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

object DataSource {
  /**
   * We have to use JVM locks because Orient refuses to lock remote databases
   */
  var wsLocks = TrieMap[WorkspaceName, StampedLock]()

  private def getLock(ws: WorkspaceName): StampedLock = {
    wsLocks.getOrElseUpdate(ws, new StampedLock())
  }

  def apply(url: String, user: String, password: String, minPoolSize: Int, maxPoolSize: Int)(implicit executionContext: ExecutionContext) = {
    val factory: OrientGraphFactory = createFactory(url, user, password)
// db pooling disabled: https://broadinstitute.atlassian.net/browse/DSDEEPB-1589
//    factory.setupPool(minPoolSize, maxPoolSize)
    new DataSource(factory)
  }

  def apply(url: String, user: String, password: String)(implicit executionContext: ExecutionContext) = {
    val factory: OrientGraphFactory = createFactory(url, user, password)
    new DataSource(factory)
  }

  def createFactory(url: String, user: String, password: String): OrientGraphFactory = {
    val factory = new OrientGraphFactory(url, user, password)
    factory.setThreadMode(THREAD_MODE.MANUAL)
    factory.setAutoStartTx(false)
    factory.setUseClassForEdgeLabel(false)
    factory.setUseLightweightEdges(true)
    factory
  }
}

class DataSource(graphFactory: OrientGraphFactory)(implicit executionContext: ExecutionContext) extends LazyLogging {

  private def lockTxn(txn: RawlsTransaction): Unit = {
    txn.readLocks.foreach { readWs =>
      DataSource.getLock(readWs).readLock()
    }
    txn.writeLocks.foreach { writeWs =>
      DataSource.getLock(writeWs).writeLock()
    }
  }

  private def unlockTxn(txn: RawlsTransaction): Unit = {
    txn.writeLocks.foreach { writeWs =>
      DataSource.getLock(writeWs).tryUnlockWrite()
    }
    txn.readLocks.foreach { readWs =>
      DataSource.getLock(readWs).tryUnlockRead()
    }
  }

  def inTransaction[T](readLocks: Set[WorkspaceName] = Set.empty[WorkspaceName],
                       writeLocks: Set[WorkspaceName] = Set.empty[WorkspaceName])(f: RawlsTransaction => T): T = {
    val graph = graphFactory.getTx
    graph.begin()
    val txn: RawlsTransaction = new RawlsTransaction(graph, this, readLocks -- writeLocks, writeLocks )
    try {
      lockTxn(txn)
      val result = f(txn)
      completeTransaction(txn)
      result
    } catch {
      case t: Throwable =>
        graph.rollback()
        throw t
    } finally {
      graph.shutdown()
      unlockTxn(txn)
    }
  }

  /**
   * This function completes or rolls back the transaction in a future at the end of the future resulting from f
   * @param f
   * @tparam T
   */
  def inFutureTransaction[T](readLocks: Set[WorkspaceName] = Set.empty[WorkspaceName],
                             writeLocks: Set[WorkspaceName] = Set.empty[WorkspaceName])(f: RawlsTransaction => Future[T]): Future[T] = {
    val graph = graphFactory.getTx
    graph.begin()
    val txn: RawlsTransaction = new RawlsTransaction(graph, this, readLocks -- writeLocks, writeLocks)
    lockTxn(txn)
    val resultFuture = f(txn)

    resultFuture.transform( { result =>
      graph.makeActive()
      try {
        completeTransaction(txn)
      } catch {
        case t: OConcurrentModificationException =>
          import scala.collection.JavaConversions._
          val v = graph.getVertex(t.getRid)
          val props = v.getPropertyKeys.map(k => k -> v.getProperty(k)).toMap
          throw new RawlsException(s"concurrent modification exception modifying ${props}", t)
        case t: Throwable =>
          throw t
      } finally {
        unlockTxn(txn)
      }
      result
    }, { throwable =>
      graph.makeActive()
      completeTransactionOnException(graph, txn)
      throwable
    })
  }

  def completeTransaction[T](rawlsTransaction: RawlsTransaction): Unit = {
    try {
      if (rawlsTransaction.isRollbackOnly) {
        logger.debug("rolling back transaction marked as rollback only")
        rawlsTransaction.graph.rollback()
      } else {
        rawlsTransaction.graph.commit()
      }
    } finally {
      rawlsTransaction.graph.shutdown()
    }
  }

  def completeTransactionOnException(graph: OrientGraph, rawlsTransaction: RawlsTransaction): Unit = {
    try {
      graph.rollback()
    } finally {
      graph.shutdown()
      unlockTxn(rawlsTransaction)
    }
  }

  def shutdown() = graphFactory.close()
}

class RawlsTransaction(val graph: OrientGraph,
                       dataSource: DataSource,
                       val readLocks: Set[WorkspaceName],
                       val writeLocks: Set[WorkspaceName])(implicit executionContext: ExecutionContext) {
  private val rollbackOnly = new AtomicBoolean(false)

  def withGraph[T](f: Graph => T) = {
    // because transactions are spanning threads we need to make sure the graph is active
    graph.makeActive()
    f(graph)
  }

  /**
   * Allows code running with a connection to rollback the transaction without throwing an exception.
   * All following modifications will be rolled back as well.
   */
  def setRollbackOnly(): Unit = {
    rollbackOnly.set(true)
  }

  def isRollbackOnly = rollbackOnly.get()
}
<commit_msg>add timeouts to waiting for obamalocks<commit_after>package org.broadinstitute.dsde.rawls.dataaccess

import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.{AtomicInteger, AtomicBoolean}
import java.util.concurrent.locks.{Lock, StampedLock, ReadWriteLock, ReentrantReadWriteLock}

import com.tinkerpop.blueprints.Element
import com.tinkerpop.blueprints.impls.orient.OrientElement
import com.orientechnologies.orient.core.exception.OConcurrentModificationException
import com.tinkerpop.blueprints.Graph
import com.tinkerpop.blueprints.impls.orient.OrientConfigurableGraph.THREAD_MODE
import com.tinkerpop.blueprints.impls.orient.{OrientGraph, OrientGraphFactory}
import com.typesafe.scalalogging.slf4j.LazyLogging
import org.broadinstitute.dsde.rawls.RawlsException
import org.broadinstitute.dsde.rawls.model.WorkspaceName

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

object DataSource {
  /**
   * We have to use JVM locks because Orient refuses to lock remote databases
   */
  var wsLocks = TrieMap[WorkspaceName, StampedLock]()

  private def getLock(ws: WorkspaceName): StampedLock = {
    wsLocks.getOrElseUpdate(ws, new StampedLock())
  }

  def apply(url: String, user: String, password: String, minPoolSize: Int, maxPoolSize: Int)(implicit executionContext: ExecutionContext) = {
    val factory: OrientGraphFactory = createFactory(url, user, password)
// db pooling disabled: https://broadinstitute.atlassian.net/browse/DSDEEPB-1589
//    factory.setupPool(minPoolSize, maxPoolSize)
    new DataSource(factory)
  }

  def apply(url: String, user: String, password: String)(implicit executionContext: ExecutionContext) = {
    val factory: OrientGraphFactory = createFactory(url, user, password)
    new DataSource(factory)
  }

  def createFactory(url: String, user: String, password: String): OrientGraphFactory = {
    val factory = new OrientGraphFactory(url, user, password)
    factory.setThreadMode(THREAD_MODE.MANUAL)
    factory.setAutoStartTx(false)
    factory.setUseClassForEdgeLabel(false)
    factory.setUseLightweightEdges(true)
    factory
  }
}

class DataSource(graphFactory: OrientGraphFactory)(implicit executionContext: ExecutionContext) extends LazyLogging {

  private def lockTxn(txn: RawlsTransaction): Unit = {
    txn.readLocks.foreach { readWs =>
      if (DataSource.getLock(readWs).tryReadLock(5, TimeUnit.SECONDS) == 0) {
        throw new RawlsException(s"Could not acquire read lock for $readWs")
      }
    }
    txn.writeLocks.foreach { writeWs =>
      if (DataSource.getLock(writeWs).tryWriteLock(30, TimeUnit.SECONDS) == 0) {
        throw new RawlsException(s"Could not acquire write lock for $writeWs")
      }
    }
  }

  private def unlockTxn(txn: RawlsTransaction): Unit = {
    txn.writeLocks.foreach { writeWs =>
      DataSource.getLock(writeWs).tryUnlockWrite()
    }
    txn.readLocks.foreach { readWs =>
      DataSource.getLock(readWs).tryUnlockRead()
    }
  }

  def inTransaction[T](readLocks: Set[WorkspaceName] = Set.empty[WorkspaceName],
                       writeLocks: Set[WorkspaceName] = Set.empty[WorkspaceName])(f: RawlsTransaction => T): T = {
    val graph = graphFactory.getTx
    graph.begin()
    val txn: RawlsTransaction = new RawlsTransaction(graph, this, readLocks -- writeLocks, writeLocks )
    try {
      lockTxn(txn)
      val result = f(txn)
      completeTransaction(txn)
      result
    } catch {
      case t: Throwable =>
        graph.rollback()
        throw t
    } finally {
      graph.shutdown()
      unlockTxn(txn)
    }
  }

  /**
   * This function completes or rolls back the transaction in a future at the end of the future resulting from f
   * @param f
   * @tparam T
   */
  def inFutureTransaction[T](readLocks: Set[WorkspaceName] = Set.empty[WorkspaceName],
                             writeLocks: Set[WorkspaceName] = Set.empty[WorkspaceName])(f: RawlsTransaction => Future[T]): Future[T] = {
    val graph = graphFactory.getTx
    graph.begin()
    val txn: RawlsTransaction = new RawlsTransaction(graph, this, readLocks -- writeLocks, writeLocks)
    lockTxn(txn)
    val resultFuture = f(txn)

    resultFuture.transform( { result =>
      graph.makeActive()
      try {
        completeTransaction(txn)
      } catch {
        case t: OConcurrentModificationException =>
          import scala.collection.JavaConversions._
          val v = graph.getVertex(t.getRid)
          val props = v.getPropertyKeys.map(k => k -> v.getProperty(k)).toMap
          throw new RawlsException(s"concurrent modification exception modifying ${props}", t)
        case t: Throwable =>
          throw t
      } finally {
        unlockTxn(txn)
      }
      result
    }, { throwable =>
      graph.makeActive()
      completeTransactionOnException(graph, txn)
      throwable
    })
  }

  def completeTransaction[T](rawlsTransaction: RawlsTransaction): Unit = {
    try {
      if (rawlsTransaction.isRollbackOnly) {
        logger.debug("rolling back transaction marked as rollback only")
        rawlsTransaction.graph.rollback()
      } else {
        rawlsTransaction.graph.commit()
      }
    } finally {
      rawlsTransaction.graph.shutdown()
    }
  }

  def completeTransactionOnException(graph: OrientGraph, rawlsTransaction: RawlsTransaction): Unit = {
    try {
      graph.rollback()
    } finally {
      graph.shutdown()
      unlockTxn(rawlsTransaction)
    }
  }

  def shutdown() = graphFactory.close()
}

class RawlsTransaction(val graph: OrientGraph,
                       dataSource: DataSource,
                       val readLocks: Set[WorkspaceName],
                       val writeLocks: Set[WorkspaceName])(implicit executionContext: ExecutionContext) {
  private val rollbackOnly = new AtomicBoolean(false)

  def withGraph[T](f: Graph => T) = {
    // because transactions are spanning threads we need to make sure the graph is active
    graph.makeActive()
    f(graph)
  }

  /**
   * Allows code running with a connection to rollback the transaction without throwing an exception.
   * All following modifications will be rolled back as well.
   */
  def setRollbackOnly(): Unit = {
    rollbackOnly.set(true)
  }

  def isRollbackOnly = rollbackOnly.get()
}

{'fraction_non_alphanumeric': 0.09724380469452851, 'fraction_numerical': 0.0012267931626727734, 'mean_word_length': 3.409664623151821, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 31, 'https://': 2, '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': '1892331', 'n_tokens_mistral': 3695, 'n_tokens_neox': 3601, 'n_words': 803}