finagle-base-http: add MapBackedHeaderMap

finagle-base-http/src/main/scala/com/twitter/finagle/http/headers/DefaultHeaderMap.scala

@@ -3,5 +3,5 @@ package com.twitter.finagle.http.headers
 import com.twitter.finagle.http.HeaderMap
 
 object DefaultHeaderMap {
-  def apply(headers: (String, String)*): HeaderMap = HashBackedHeaderMap(headers: _*)
+  def apply(headers: (String, String)*): HeaderMap = JTreeMapBackedHeaderMap(headers: _*)
 }

finagle-base-http/src/main/scala/com/twitter/finagle/http/headers/HashBackedHeaderMap.scala

@@ -82,7 +82,7 @@ private final class HashBackedHeaderMap extends http.HeaderMap {
     }
 }
 
-private object HashBackedHeaderMap {
+object HashBackedHeaderMap {
 
   /** Construct a new `HeaderMap` with the header list
    *

finagle-base-http/src/main/scala/com/twitter/finagle/http/headers/Header.scala

@@ -1,7 +1,9 @@
 package com.twitter.finagle.http.headers
 
-import com.twitter.finagle.http.HeaderMap
+import scala.annotation.tailrec
 import scala.collection.mutable
+import scala.collection.AbstractIterator
+import com.twitter.finagle.http.HeaderMap
 
 private[http] sealed class Header private (final val name: String, final val value: String)
     extends HeaderMap.NameValue {
@@ -14,6 +16,21 @@ private[http] sealed class Header private (final val name: String, final val val
 private[http] object Header {
 
   final class Root private[Header] (name: String, value: String) extends Header(name, value) {
+    self => 
+
+    def iterator: Iterator[HeaderMap.NameValue] =
+      if (next == null) Iterator.single(this)
+      else {
+        new AbstractIterator[HeaderMap.NameValue] {
+          var cur: Header = self
+          def hasNext: Boolean = cur != null
+          def next(): HeaderMap.NameValue = {
+            var n = cur
+            cur = n.next
+            n
+          }
+        }
+      }
 
     // We want to keep a reference to the last element of this linked list
     // so we don't need to traverse the whole list to add an element.
@@ -49,4 +66,65 @@ private[http] object Header {
   /** Create a new root node */
   def root(name: String, value: String): Root =
     new Root(name, value)
+
+  def uniqueNames(orig: Iterator[Header]): Iterator[String] = new Iterator[String] {
+    private[this] val it = orig
+    private[this] var current: Iterator[String] = Iterator.empty
+
+    private[this] def collectUnique(from: Header): Iterator[String] = new Iterator[String] {
+
+      // We need to keep `_next` current so that we can reliably determine `.hasNext`.
+      // That means that if `_next` is not null, it is, in fact, the next value.
+      private[this] var nextUnique = from
+
+      // We keep a set of observed names so that we don't accidentally return
+      // duplicates. Another way to have done this would be to simply add all names
+      // to a set and return an iterator over that set, but this way is more lazy.
+      private[this] var seen = Set.empty + from.name
+
+      def hasNext: Boolean = nextUnique != null
+
+      def next(): String = {
+        if (!hasNext) throw new NoSuchElementException
+
+        val result = nextUnique.name
+        // We then need to determine the next `_next` so we can
+        // again have a sane `.hasNext`.
+        prepareNext()
+        result
+      }
+
+      @tailrec
+      private[this] def prepareNext(): Unit = {
+        nextUnique = nextUnique.next
+        if (nextUnique != null) {
+          val prevSize = seen.size
+          seen = seen + nextUnique.name
+          if (seen.size == prevSize) {
+            // already observed. Try again.
+            prepareNext()
+          }
+        }
+      }
+    }
+
+    def hasNext: Boolean =
+      it.hasNext || current.hasNext
+
+    def next(): String = {
+      if (current.isEmpty) {
+        val hs = it.next()
+        if (hs.next == null) {
+          // A shortcut to see if we have only a single entry for this name
+          hs.name
+        } else {
+          // slow path: we need to make a sub-iterator
+          current = collectUnique(hs)
+          current.next()
+        }
+      } else {
+        current.next()
+      }
+    }
+  }
 }

finagle-base-http/src/main/scala/com/twitter/finagle/http/headers/HeadersHash.scala

@@ -9,10 +9,7 @@ import scala.collection.mutable
 // - iteration by gaining access to `entriesIterator` (protected method).
 // - get/add functions by providing custom hashCode and equals methods for a key
 private[http] final class HeadersHash extends mutable.HashMap[String, Header.Root] {
-
-  private def hashChar(c: Char): Int =
-    if (c >= 'A' && c <= 'Z') c + 32
-    else c
+  import HeadersHash._
 
   // Adopted from Netty 3 HttpHeaders.
   override protected def elemHashCode(key: String): Int = {
@@ -68,67 +65,9 @@ private[http] final class HeadersHash extends mutable.HashMap[String, Header.Roo
     array.iterator
   }
 
-  def uniqueNamesIterator: Iterator[String] = new Iterator[String] {
-    private[this] val it = copiedEntitiesIterator
-    private[this] var current: Iterator[String] = Iterator.empty
-
-    private[this] def collectUnique(from: Header): Iterator[String] = new Iterator[String] {
-
-      // We need to keep `_next` current so that we can reliably determine `.hasNext`.
-      // That means that if `_next` is not null, it is, in fact, the next value.
-      private[this] var nextUnique = from
-
-      // We keep a set of observed names so that we don't accidentally return
-      // duplicates. Another way to have done this would be to simply add all names
-      // to a set and return an iterator over that set, but this way is more lazy.
-      private[this] var seen = Set.empty + from.name
-
-      def hasNext: Boolean = nextUnique != null
-
-      def next(): String = {
-        if (!hasNext) throw new NoSuchElementException
-
-        val result = nextUnique.name
-        // We then need to determine the next `_next` so we can
-        // again have a sane `.hasNext`.
-        prepareNext()
-        result
-      }
-
-      @tailrec
-      private[this] def prepareNext(): Unit = {
-        nextUnique = nextUnique.next
-        if (nextUnique != null) {
-          val prevSize = seen.size
-          seen = seen + nextUnique.name
-          if (seen.size == prevSize) {
-            // already observed. Try again.
-            prepareNext()
-          }
-        }
-      }
-    }
-
-    def hasNext: Boolean =
-      it.hasNext || current.hasNext
-
-    def next(): String = {
-      if (current.isEmpty) {
-        val hs = it.next()
-        if (hs.next == null) {
-          // A shortcut to see if we have only a single entry for this name
-          hs.name
-        } else {
-          // slow path: we need to make a sub-iterator
-          current = collectUnique(hs)
-          current.next()
-        }
-      } else {
-        current.next()
-      }
-    }
-  }
-
+  def uniqueNamesIterator: Iterator[String] =
+    Header.uniqueNames(copiedEntitiesIterator)
+
   def flattenIterator: Iterator[(String, String)] =
     flattenedNameValueIterator.map(nv => (nv.name, nv.value))
 
@@ -181,3 +120,9 @@ private[http] final class HeadersHash extends mutable.HashMap[String, Header.Roo
   def removeAll(key: String): Unit = remove(key)
 
 }
+
+object HeadersHash {
+  final def hashChar(c: Char): Int =
+    if (c >= 'A' && c <= 'Z') c + 32
+    else c
+}

finagle-base-http/src/main/scala/com/twitter/finagle/http/headers/JTreeMapBackedHeaderMap.scala

@@ -0,0 +1,139 @@
+package com.twitter.finagle.http.headers
+
+import com.twitter.finagle.http.HeaderMap
+import scala.annotation.tailrec
+import java.util.function.BiConsumer
+import scala.collection.AbstractIterator
+
+/**
+ * Mutable, thread-safe [[HeaderMap]] implementation, backed by
+ * a mutable [[Map[String, Header]]], where the map key
+ * is forced to lower case
+ */
+final private class JTreeMapBackedHeaderMap extends HeaderMap {
+  import HeaderMap._
+
+  // In general, Map's that are not thread safe are not
+  // durable to concurrent modification and can result in infinite loops
+  // and exceptions.
+  // As such, we synchronize on the underlying collection when performing
+  // accesses to avoid this. In the common case of no concurrent access,
+  // this should be cheap.
+  val underlying: java.util.TreeMap[String, Header.Root] = 
+  new java.util.TreeMap[String, Header.Root](JTreeMapBackedHeaderMap.SharedComparitor)
+
+  private def foreachConsumer[U](f: ((String, String)) => U):
+    BiConsumer[String, Header.Root] = new BiConsumer[String, Header.Root](){
+      def accept(key: String, header: Header.Root): Unit = header.iterator.foreach(
+        nv => f(nv.name, nv.value)
+      )
+    }
+
+  override def foreach[U](f: ((String, String)) => U): Unit = 
+    underlying.forEach(foreachConsumer(f))
+
+  // ---- HeaderMap -----
+
+  // Validates key and value.
+  def add(key: String, value: String): this.type = {
+    validateName(key)
+    addUnsafe(key, foldReplacingValidateValue(key, value))
+  }
+
+  // Validates key and value.
+  def set(key: String, value: String): this.type = {
+    validateName(key)
+    setUnsafe(key, foldReplacingValidateValue(key, value))
+  }
+
+  // ---- Map/MapLike -----
+
+  def -=(key: String): this.type = removed(key)
+  def +=(kv: (String, String)): this.type = set(kv._1, kv._2)
+
+  /**
+   * Underlying headers eagerly copied to an array, without synchronizing
+   * on the underlying collection.
+   */
+  private[this] def copyHeaders: Iterator[Header.Root] = underlying.values.toArray(new Array[Header.Root](underlying.size)).iterator
+
+  def iterator: Iterator[(String, String)] =
+    nameValueIterator.map(nv => (nv.name, nv.value))
+
+  override def keys: Set[String] = keysIterator.toSet
+
+  override def keysIterator: Iterator[String] = underlying.synchronized {
+    Header.uniqueNames(copyHeaders)
+  }
+
+  private[finagle] final override def nameValueIterator: Iterator[HeaderMap.NameValue] =
+    underlying.synchronized {
+      copyHeaders.flatMap(_.iterator)
+    }
+
+  def getAll(key: String): Seq[String] = underlying.synchronized {
+    underlying.get(key) match {
+      case null => Nil
+      case r: Header.Root => r.values
+    }
+  }
+
+  // Does not validate key and value.
+  def addUnsafe(key: String, value: String): this.type  = underlying.synchronized {
+    underlying.get(key) match {
+      case null => underlying.put(key, Header.root(key, value))
+      case h    => h.add(key, value)
+    }
+    this
+  }
+
+  // Does not validate key and value.
+  def setUnsafe(key: String, value: String): this.type = underlying.synchronized {
+    underlying.put(key, Header.root(key, value))
+    this
+  }
+
+  def get(key: String): Option[String] = underlying.synchronized {
+    underlying.get(key) match {
+      case null => None
+      case h => Some(h.value)
+    }
+  }
+
+  def removed(key: String): this.type = underlying.synchronized {
+    underlying.remove(key)
+    this
+  }
+}
+
+
+object JTreeMapBackedHeaderMap {
+
+  val SharedComparitor = new java.util.Comparator[String] {
+    def compare(key1: String, key2: String): Int = {
+      // Shorter strings are always less, regardless of their content
+      val lenthDiff = key1.length - key2.length
+      if (lenthDiff != 0) lenthDiff
+      else {
+        @tailrec
+        def go(i: Int): Int = {
+          if (i == key1.length) 0 // end, they are equal.
+          else {
+            val char1 = HeadersHash.hashChar(key1.charAt(i))
+            val char2 = HeadersHash.hashChar(key2.charAt(i))
+            val diff = char1 - char2
+            if (diff == 0) go(i + 1)
+            else diff
+          }
+        }
+        go(0)
+      }
+    }
+  }
+
+  def apply(headers: (String, String)*): HeaderMap = {
+    val result = new JTreeMapBackedHeaderMap
+    headers.foreach(t => result.add(t._1, t._2))
+    result
+  }
+}

finagle-base-http/src/test/scala/com/twitter/finagle/http/headers/DefaultHeaderMapTest.scala

@@ -2,6 +2,10 @@ package com.twitter.finagle.http.headers
 
 import com.twitter.finagle.http.HeaderMap
 
-class DefaultHeaderMapTest extends AbstractHeaderMapTest {
-  final def newHeaderMap(headers: (String, String)*): HeaderMap = DefaultHeaderMap(headers: _*)
+class HashBackedHeaderMapTest extends AbstractHeaderMapTest {
+  final def newHeaderMap(headers: (String, String)*): HeaderMap = HashBackedHeaderMap(headers: _*)
 }
+
+class JTreeMapBackedHeaderMapTest extends AbstractHeaderMapTest {
+  final def newHeaderMap(headers: (String, String)*): HeaderMap = JTreeMapBackedHeaderMap(headers: _*)
+}