swaydb.core.segment.SegmentAssigner.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of core_2.12 Show documentation
core
The newest version!
/*
 * Copyright (c) 2020 Simer JS Plaha ([email protected] - @simerplaha)
 *
 * This file is a part of SwayDB.
 *
 * SwayDB is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * SwayDB is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with SwayDB. If not, see .
 *
 * Additional permission under the GNU Affero GPL version 3 section 7:
 * If you modify this Program or any covered work, only by linking or combining
 * it with separate works, the licensors of this Program grant you additional
 * permission to convey the resulting work.
 */

package swaydb.core.segment

import swaydb.core.data.{KeyValue, Memory, MemoryOption, Value}
import swaydb.core.util.DropIterator
import swaydb.core.util.skiplist.SkipList
import swaydb.data.order.KeyOrder
import swaydb.data.slice.{Slice, SliceOption}

import scala.annotation.tailrec
import scala.collection.mutable

private[core] object SegmentAssigner {

  def assignMinMaxOnlyUnsafe(inputSegments: Iterable[Segment],
                             targetSegments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): Iterable[Segment] =
    SegmentAssigner.assignUnsafe(2 * inputSegments.size, Segment.tempMinMaxKeyValues(inputSegments), targetSegments).keys

  def assignMinMaxOnlyUnsafe(keyValues: Either[SkipList[SliceOption[Byte], MemoryOption, Slice[Byte], Memory], Slice[Memory]],
                             targetSegments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): Iterable[Segment] =
    keyValues match {
      case Left(value) =>
        assignMinMaxOnlyUnsafe(value, targetSegments)

      case Right(value) =>
        assignMinMaxOnlyUnsafe(value, targetSegments)
    }

  def assignMinMaxOnlyUnsafe(input: SkipList[SliceOption[Byte], MemoryOption, Slice[Byte], Memory],
                             targetSegments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): Iterable[Segment] =
    SegmentAssigner.assignUnsafe(2, Segment.tempMinMaxKeyValues(input), targetSegments).keys

  def assignMinMaxOnlyUnsafe(input: Slice[Memory],
                             targetSegments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): Iterable[Segment] =
    SegmentAssigner.assignUnsafe(2, Segment.tempMinMaxKeyValues(input), targetSegments).keys

  def assignUnsafe(keyValues: Slice[KeyValue],
                   segments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): mutable.Map[Segment, Slice[KeyValue]] =
    assignUnsafe(
      keyValuesCount = keyValues.size,
      keyValues = keyValues,
      segments = segments
    )

  /**
   * @param keyValuesCount keyValuesCount is needed here because keyValues could be a [[ConcurrentSkipList]]
   *                       where calculating size is not constant time.
   */
  def assignUnsafe(keyValuesCount: Int,
                   keyValues: Iterable[KeyValue],
                   segments: Iterable[Segment])(implicit keyOrder: KeyOrder[Slice[Byte]]): mutable.Map[Segment, Slice[KeyValue]] = {
    if (Segment hasOnlyOneSegment segments) { //.size iterates the entire Iterable which is not needed.
      mutable.Map((segments.head, Slice.from(keyValues, keyValuesCount)))
    } else {
      import keyOrder._

      val assignmentsMap = mutable.Map.empty[Segment, Slice[KeyValue]]
      val segmentsIterator = segments.iterator

      def getNextSegmentMayBe() = if (segmentsIterator.hasNext) segmentsIterator.next() else Segment.Null

      def assignKeyValueToSegment(segment: Segment,
                                  keyValue: KeyValue,
                                  remainingKeyValues: Int): Unit =
        assignmentsMap.get(segment) match {
          case Some(currentAssignments) =>
            currentAssignments add keyValue

          case None =>
            //+1 for cases when a Range might extend to the next Segment.
            val initial = Slice.of[KeyValue](remainingKeyValues + 1)
            initial add keyValue
            assignmentsMap += (segment -> initial)
        }

      @tailrec
      def assign(remainingKeyValues: DropIterator[Memory.Range, KeyValue],
                 thisSegmentMayBe: SegmentOption,
                 nextSegmentMayBe: SegmentOption): Unit =
        (remainingKeyValues.headOrNull, thisSegmentMayBe, nextSegmentMayBe) match {
          //add this key-value if it is the new smallest key or if this key belong to this Segment or if there is no next Segment
          case (keyValue: KeyValue, thisSegment: Segment, _) if keyValue.key <= thisSegment.minKey || Segment.belongsTo(keyValue, thisSegment) || nextSegmentMayBe.isNoneS =>
            keyValue match {
              case keyValue: KeyValue.Fixed =>
                assignKeyValueToSegment(thisSegment, keyValue, remainingKeyValues.size)
                assign(remainingKeyValues.dropHead(), thisSegmentMayBe, nextSegmentMayBe)

              case keyValue: KeyValue.Range =>
                nextSegmentMayBe match {
                  case nextSegment: Segment if keyValue.toKey > nextSegment.minKey =>
                    val (fromValue, rangeValue) = keyValue.fetchFromAndRangeValueUnsafe
                    val thisSegmentsRange = Memory.Range(fromKey = keyValue.fromKey, toKey = nextSegment.minKey, fromValue = fromValue, rangeValue = rangeValue)
                    val nextSegmentsRange = Memory.Range(fromKey = nextSegment.minKey, toKey = keyValue.toKey, fromValue = Value.FromValue.Null, rangeValue = rangeValue)

                    assignKeyValueToSegment(thisSegment, thisSegmentsRange, remainingKeyValues.size)
                    assign(remainingKeyValues.dropPrepend(nextSegmentsRange), nextSegment, getNextSegmentMayBe())

                  case _ =>
                    assignKeyValueToSegment(thisSegment, keyValue, remainingKeyValues.size)
                    assign(remainingKeyValues.dropHead(), thisSegmentMayBe, nextSegmentMayBe)
                }
            }


          // is this a gap key between thisSegment and the nextSegment
          case (keyValue: KeyValue, thisSegment: Segment, nextSegment: Segment) if keyValue.key < nextSegment.minKey =>
            keyValue match {
              case keyValue: KeyValue.Fixed =>
                //ignore if a key-value is not already assigned to thisSegment. No point adding a single key-value to a Segment.
                if (assignmentsMap.contains(thisSegment)) {
                  assignKeyValueToSegment(thisSegment, keyValue, remainingKeyValues.size)
                  assign(remainingKeyValues.dropHead(), thisSegmentMayBe, nextSegmentMayBe)
                } else {
                  assign(remainingKeyValues, nextSegment, getNextSegmentMayBe())
                }

              case keyValue: KeyValue.Range =>
                nextSegmentMayBe match {
                  case nextSegment: Segment if keyValue.toKey > nextSegment.minKey =>
                    //if it's a gap Range key-value and it's flows onto the next Segment, just jump to the next Segment.
                    assign(remainingKeyValues, nextSegment, getNextSegmentMayBe())

                  case _ =>
                    //ignore if a key-value is not already assigned to thisSegment. No point adding a single key-value to a Segment.
                    //same code as above, need to push it to a common function.
                    if (assignmentsMap.contains(thisSegment)) {
                      assignKeyValueToSegment(thisSegment, keyValue, remainingKeyValues.size)
                      assign(remainingKeyValues.dropHead(), thisSegmentMayBe, nextSegmentMayBe)
                    } else {
                      assign(remainingKeyValues, nextSegment, getNextSegmentMayBe())
                    }
                }
            }

          case (_: KeyValue, _: Segment, nextSegment: Segment) =>
            assign(remainingKeyValues, nextSegment, getNextSegmentMayBe())

          case (_, _, _) =>
            ()
        }

      if (segmentsIterator.hasNext) {
        assign(DropIterator(keyValuesCount, keyValues.iterator), segmentsIterator.next(), getNextSegmentMayBe())
        assignmentsMap map {
          case (segment, keyValues) =>
            (segment, keyValues.close())
        }
      } else {
        assignmentsMap
      }
    }
  }
}