使用Akka流动态压缩列表[source[ByteString，NotUsed]]

Question

我有一个source[ByteString，NotUsed]列表，它与来自S3 bucket的文件名成对。这些需要压缩在恒定内存中，并在Play2.6中提供服务。

我的变化主要在这里：

override def onPush(): Unit = {
  val (filepath, source: StreamGenerator) = grab(in)
  buffer.startEntry(filepath)
  val src: Source[ByteString, NotUsed] = source()
  val operation = src.runForeach(bytestring => {
    val byteInputStream = new ByteArrayInputStream(bytestring.toArray)
    emitMultiple(out, fileChunks(byteInputStream, buffer))
  })
  operation.onComplete {
    case _ => buffer.endEntry()
  }
  Await.ready(operation, 5.minute)
}

我知道这是阻塞，但我不确定在这种情况下是否允许。

我如何以安全的方式完成这件事？

override def onPush(): Unit = {
  val (filepath, source: StreamGenerator) = grab(in)
  buffer.startEntry(filepath)
  val stream = source().runWith(StreamConverters.asInputStream(1.minute))
  currentStream = Some(stream)
  emitMultiple(out, fileChunks(stream, buffer), () => buffer.endEntry())
}

下面是我当前图形阶段实现的完整列表：

import java.io.{ByteArrayInputStream, InputStream, OutputStream}

import akka.NotUsed
import akka.stream._
import akka.stream.scaladsl._
import akka.stream.stage._
import akka.util.{ByteString, ByteStringBuilder}

import scala.concurrent.duration._
import scala.concurrent.{Await, ExecutionContext}
import scala.util.control.NonFatal

//scalastyle:off
class StreamedZip(bufferSize: Int = 64 * 1024)(implicit ec: ExecutionContext,
                                               mat: ActorMaterializer)
    extends GraphStage[FlowShape[StreamedZip.ZipSource, ByteString]] {

  import StreamedZip._

  val in: Inlet[ZipSource]    = Inlet("StreamedZip.in")
  val out: Outlet[ByteString] = Outlet("StreamedZip.out")
  override val shape          = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with StageLogging {
      private val buffer                             = new ZipBuffer(bufferSize)
      private var currentStream: Option[InputStream] = None

      setHandler(
        out,
        new OutHandler {
          override def onPull(): Unit =
            if (isClosed(in)) {
              if (buffer.isEmpty) completeStage()
              else {
                buffer.close
                push(out, buffer.toByteString)
              }
            } else pull(in)

          override def onDownstreamFinish(): Unit = {
            closeInput()
            buffer.close
            super.onDownstreamFinish()
          }
        }
      )

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            val (filepath, source: StreamGenerator) = grab(in)
            buffer.startEntry(filepath)
            val stream = source().runWith(StreamConverters.asInputStream(1.minute))
            emitMultiple(out, fileChunks(stream, buffer), () => { buffer.endEntry() })
          }

          override def onUpstreamFinish(): Unit = {
            println("Updstream finish")
            closeInput()
            if (buffer.isEmpty) completeStage()
            else {
              buffer.close()
              if (isAvailable(out)) {
                push(out, buffer.toByteString)
              }
            }
          }
        }
      )

      private def closeInput(): Unit = {
        currentStream.foreach(_.close)
        currentStream = None
      }

      private def fileChunks(stream: InputStream, buffer: ZipBuffer): Iterator[ByteString] = {
        // This seems like a good trade-off between single-byte
        // read I/O performance and doubling the ZipBuffer size.
        //
        // And it's still a decent defense against DDOS resource
        // limit attacks.
        val readBuffer = new Array[Byte](1024)
        var done       = false

        def result: Stream[ByteString] =
          if (done) Stream.empty
          else {
            try {
              while (!done && buffer.remaining > 0) {
                val bytesToRead = Math.min(readBuffer.length, buffer.remaining)
                val count       = stream.read(readBuffer, 0, bytesToRead)
                if (count == -1) {
                  stream.close
                  done = true
                } else buffer.write(readBuffer, count)
              }
              buffer.toByteString #:: result
            } catch {
              case NonFatal(e) =>
                closeInput()
                throw e
            }
          }

        result.iterator
      }
    }
}

object StreamedZip {
  type ZipFilePath     = String
  type StreamGenerator = () => Source[ByteString, NotUsed]
  type ZipSource       = (ZipFilePath, StreamGenerator)

  def apply()(implicit ec: ExecutionContext, mat: ActorMaterializer) = new StreamedZip()

}

class ZipBuffer(val bufferSize: Int = 64 * 1024) {
  import java.util.zip.{ZipEntry, ZipOutputStream}

  private var builder = new ByteStringBuilder()
  private val zip = new ZipOutputStream(builder.asOutputStream) {
    // this MUST ONLY be used after flush()!
    def setOut(newStream: OutputStream): Unit = out = newStream
  }
  private var inEntry = false
  private var closed  = false

  def close(): Unit = {
    endEntry()
    closed = true
    zip.close()
  }

  def remaining(): Int = bufferSize - builder.length

  def isEmpty(): Boolean = builder.isEmpty

  def startEntry(path: String): Unit =
    if (!closed) {
      endEntry()
      zip.putNextEntry(new ZipEntry(path))
      inEntry = true
    }

  def endEntry(): Unit =
    if (!closed && inEntry) {
      inEntry = false
      zip.closeEntry()
    }

  def write(byte: Int): Unit =
    if (!closed && inEntry) zip.write(byte)

  def write(bytes: Array[Byte], length: Int): Unit =
    if (!closed && inEntry) zip.write(bytes, 0, length)

  def toByteString(): ByteString = {
    zip.flush()
    val result = builder.result
    builder = new ByteStringBuilder()
    // set the underlying output for the zip stream to be the buffer
    // directly, so we don't have to copy the zip'd byte array.
    zip.setOut(builder.asOutputStream)
    result
  }
}

Answer 1

最后我使用了上面的zipbuffer，并使用akka流DSL解决了整个问题。

我的解决方案如下所示：

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.scaladsl.Source
import akka.stream.{ActorMaterializer, SourceShape}
import akka.util.ByteString
import com.company.config.AWS
import org.log4s.getLogger

case class S3StreamingServiceLike(awsConf: AWS, s3Client: S3ClientAlpakka)(
    implicit sys: ActorSystem,
    mat: ActorMaterializer)
    extends S3StreamingService {

  private implicit class ConcatSyntax[T, U](source: Source[T, U]) {
    def ++[TT >: T, NotUsed](that: Source[SourceShape[TT], NotUsed]): Source[Any, U] = //scalastyle:ignore
      source.concat(that)
  }

  private val logger = getLogger

  private sealed trait ZipElement
  private case class FileStart(name: String, index: Int, outOf: Int) extends ZipElement
  private case class FileEnd(name: String, index: Int, outOf: Int)   extends ZipElement
  private case class FilePayload(byteString: ByteString)             extends ZipElement
  private case object EndZip                                         extends ZipElement

  private def payloadSource(filename: String) =
    s3Client.download(awsConf.s3BucketName, filename).map(FilePayload.apply)

  private def fileNameToZipElements(filename: String,
                                    index: Int,
                                    outOf: Int): Source[ZipElement, NotUsed] =
    Source.single(FileStart(filename, index, outOf)) ++
      payloadSource(filename) ++
      Source.single(FileEnd(filename, index, outOf))

  def streamFilesAsZip(filenames: List[String])(forUser: String): Source[ByteString, NotUsed] = {

    val zipBuffer = new ZipBuffer()

    val zipElementSource: Source[ZipElement, NotUsed] =
      Source(filenames.zipWithIndex).flatMapConcat {
        case (filename, index) => fileNameToZipElements(filename, index + 1, filenames.length)
      } ++ Source.single(EndZip)

    zipElementSource
      .map {
        case FileStart(name, index, outOf) =>
          logger.info(s"Zipping file #$index of $outOf with name $name for user $forUser")
          zipBuffer.startEntry(name)
          None
        case FilePayload(byteString) =>
          if (byteString.length > zipBuffer.remaining()) {
            throw new Exception(
              s"Bytestring size exceeded buffer size ${byteString.length} > ${zipBuffer.remaining}")
          }
          zipBuffer.write(byteString.toArray, byteString.length)
          Some(zipBuffer.toByteString())
        case FileEnd(name, index, outOf) =>
          logger.info(s"Finished zipping file #$index of $outOf with $name for user $forUser")
          zipBuffer.endEntry()
          Some(zipBuffer.toByteString())
        case EndZip =>
          zipBuffer.close()
          Some(zipBuffer.toByteString())
      }
      .collect {
        case Some(bytes) if bytes.length > 0 => bytes
      }
  }

}