<commit_before>package com.twitter.finagle.b3.thrift

import com.twitter.util.Duration
import collection.mutable.HashMap
import com.twitter.finagle.tracing.TraceId
import com.twitter.finagle.util.Timer

/**
 * Takes care of storing the spans in a thread safe fashion. If a span
 * is not removed from the map it will expire after the deadline is reached
 * and sent off to B3 despite being incomplete.
 */
class DeadlineSpanMap(tracer: BigBrotherBirdTracer, deadline: Duration) {

  private[this] val spanMap = HashMap[TraceId, Span]()
  private[this] val timer = Timer.default

  /**
   * Update the span in the map. If none exists create a new span. Synchronized.
   */
  def update(traceId: TraceId)(f: Span => Span): Span = synchronized {
    val span = spanMap.get(traceId) match {
      case Some(s) => f(s)
      case None =>
        // no span found, let's create a new one
        val span = f(Span(traceId))

        // if this new span isn't triggered by a natural end we
        // send off what we have anyway
        timer.schedule(deadline) {
          remove(traceId) foreach { tracer.logSpan(_) }
        }

        span
    }

    spanMap.put(traceId, span)
    span
  }

  /**
   * Remove this span from the map. Synchronized.<commit_msg>[split] Only run remove code snippet once, not repeatedly.<commit_after>package com.twitter.finagle.b3.thrift

import com.twitter.util.Duration
import collection.mutable.HashMap
import com.twitter.finagle.tracing.TraceId
import com.twitter.finagle.util.Timer

/**
 * Takes care of storing the spans in a thread safe fashion. If a span
 * is not removed from the map it will expire after the deadline is reached
 * and sent off to B3 despite being incomplete.
 */
class DeadlineSpanMap(tracer: BigBrotherBirdTracer, deadline: Duration) {

  private[this] val spanMap = HashMap[TraceId, Span]()
  private[this] val timer = Timer.default

  /**
   * Update the span in the map. If none exists create a new span. Synchronized.
   */
  def update(traceId: TraceId)(f: Span => Span): Span = synchronized {
    val span = spanMap.get(traceId) match {
      case Some(s) => f(s)
      case None =>
        // no span found, let's create a new one
        val span = f(Span(traceId))

        // if this new span isn't triggered by a natural end we
        // send off what we have anyway
        timer.schedule(deadline.fromNow) {
          remove(traceId) foreach { tracer.logSpan(_) }
        }

        span
    }

    spanMap.put(traceId, span)
    span
  }

  /**
   * Remove this span from the map. Synchronized.
{'fraction_non_alphanumeric': 0.09466303077522399, 'fraction_numerical': 0.0015582391897156214, 'mean_word_length': 2.8385650224215246, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '3112439', 'n_tokens_mistral': 779, 'n_tokens_neox': 758, 'n_words': 273}