forked from bluelabsio/s3-stream
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
bluelabsio#12: Add ability to buffer chunks to disk instead of memory
This allows a (much) creater amount of upload streams to run in parallel.
- Loading branch information
Showing
13 changed files
with
303 additions
and
148 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| package com.bluelabs.s3stream | ||
|
|
||
| import akka.stream.scaladsl.Source | ||
| import akka.NotUsed | ||
| import akka.util.ByteString | ||
|
|
||
| case class Chunk(data: Source[ByteString,NotUsed], size: Int) |
51 changes: 0 additions & 51 deletions
51
s3-stream/src/main/scala/com/bluelabs/s3stream/Chunker.scala
This file was deleted.
Oops, something went wrong.
98 changes: 98 additions & 0 deletions
98
s3-stream/src/main/scala/com/bluelabs/s3stream/DiskBuffer.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| package com.bluelabs.s3stream | ||
|
|
||
| import java.io.FileOutputStream | ||
| import java.io.RandomAccessFile | ||
| import java.nio.channels.FileChannel | ||
| import java.nio.file.Files | ||
| import java.util.concurrent.atomic.AtomicInteger | ||
|
|
||
| import akka.NotUsed | ||
| import akka.stream.ActorAttributes | ||
| import akka.stream.Attributes | ||
| import akka.stream.FlowShape | ||
| import akka.stream.Inlet | ||
| import akka.stream.Outlet | ||
| import akka.stream.scaladsl.FileIO | ||
| import akka.stream.stage.GraphStage | ||
| import akka.stream.stage.GraphStageLogic | ||
| import akka.stream.stage.InHandler | ||
| import akka.stream.stage.OutHandler | ||
| import akka.util.ByteString | ||
| import java.nio.file.Path | ||
|
|
||
| /** | ||
| * Buffers the complete incoming stream into a file, which can then be read several times afterwards. | ||
| * | ||
| * The stage waits for the incoming stream to complete. After that, it emits a single Chunk item on its output. The Chunk | ||
| * contains a bytestream source that can be materialized multiple times, and the total size of the file. | ||
| * | ||
| * @param maxMaterializations Number of expected materializations for the completed chunk. After this, the temp file is deleted. | ||
| * @param maximumSize Maximum size on disk to buffer | ||
| */ | ||
| class DiskBuffer (maxMaterializations: Int, maxSize: Int, tempPath: Option[Path]) extends GraphStage[FlowShape[ByteString, Chunk]] { | ||
| if (maxMaterializations < 1) throw new IllegalArgumentException("maxMaterializations should be at least 1") | ||
| if (maxSize < 1) throw new IllegalArgumentException("maximumSize should be at least 1") | ||
|
|
||
| val in = Inlet[ByteString]("in") | ||
| val out = Outlet[Chunk]("out") | ||
| override val shape = FlowShape.of(in, out) | ||
|
|
||
| override def initialAttributes = ActorAttributes.dispatcher("akka.stream.default-blocking-io-dispatcher") | ||
|
|
||
| override def createLogic(attr: Attributes): GraphStageLogic = new GraphStageLogic(shape) { | ||
| val path:Path = tempPath.map(dir => Files.createTempFile(dir, "s3-buffer-", ".bin")).getOrElse(Files.createTempFile("s3-buffer-", ".bin")) | ||
| path.toFile().deleteOnExit() | ||
| val writeBuffer = new RandomAccessFile(path.toFile(), "rw").getChannel().map(FileChannel.MapMode.READ_WRITE, 0, maxSize) | ||
| var length = 0 | ||
|
|
||
| setHandler(out, new OutHandler { | ||
| override def onPull(): Unit = { | ||
| if (isClosed(in)) { | ||
| emit() | ||
| } else { | ||
| pull(in) | ||
| } | ||
| } | ||
| }) | ||
|
|
||
| setHandler(in, new InHandler { | ||
| override def onPush(): Unit = { | ||
| val elem = grab(in) | ||
| length += elem.size | ||
| writeBuffer.put(elem.asByteBuffer) | ||
| pull(in) | ||
| } | ||
|
|
||
| override def onUpstreamFinish(): Unit = { | ||
| if (isAvailable(out)) { | ||
| emit() | ||
| } | ||
| completeStage() | ||
| } | ||
| }) | ||
|
|
||
| private def emit(): Unit = { | ||
| // TODO Should we do http://stackoverflow.com/questions/2972986/how-to-unmap-a-file-from-memory-mapped-using-filechannel-in-java ? | ||
| writeBuffer.force() | ||
|
|
||
| val ch = new FileOutputStream(path.toFile(), true).getChannel() | ||
| try { | ||
| ch.truncate(length) | ||
| } finally { | ||
| ch.close() | ||
| } | ||
|
|
||
| val deleteCounter = new AtomicInteger(maxMaterializations) | ||
| val src = FileIO.fromPath(path, 65536).mapMaterializedValue { f => | ||
| if (deleteCounter.decrementAndGet() <= 0) { | ||
| import scala.concurrent.ExecutionContext.Implicits.global | ||
| f.onComplete { _ => | ||
| path.toFile().delete() | ||
| } | ||
| } | ||
| NotUsed | ||
| } | ||
| emit(out, Chunk(src, length), () => completeStage()) | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
69 changes: 69 additions & 0 deletions
69
s3-stream/src/main/scala/com/bluelabs/s3stream/MemoryBuffer.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| package com.bluelabs.s3stream | ||
|
|
||
| import java.io.FileOutputStream | ||
| import java.io.RandomAccessFile | ||
| import java.nio.channels.FileChannel | ||
| import java.nio.file.Files | ||
| import java.util.concurrent.atomic.AtomicInteger | ||
|
|
||
| import akka.NotUsed | ||
| import akka.stream.ActorAttributes | ||
| import akka.stream.Attributes | ||
| import akka.stream.FlowShape | ||
| import akka.stream.Inlet | ||
| import akka.stream.Outlet | ||
| import akka.stream.scaladsl.FileIO | ||
| import akka.stream.stage.GraphStage | ||
| import akka.stream.stage.GraphStageLogic | ||
| import akka.stream.stage.InHandler | ||
| import akka.stream.stage.OutHandler | ||
| import akka.util.ByteString | ||
| import akka.stream.scaladsl.Source | ||
|
|
||
| /** | ||
| * Buffers the complete incoming stream into memory, which can then be read several times afterwards. | ||
| * | ||
| * The stage waits for the incoming stream to complete. After that, it emits a single Chunk item on its output. The Chunk | ||
| * contains a bytestream source that can be materialized multiple times, and the total size of the file. | ||
| * | ||
| * @param maximumSize Maximum size to buffer | ||
| */ | ||
| class MemoryBuffer(maxSize: Int) extends GraphStage[FlowShape[ByteString, Chunk]] { | ||
| val in = Inlet[ByteString]("in") | ||
| val out = Outlet[Chunk]("out") | ||
| override val shape = FlowShape.of(in, out) | ||
|
|
||
| override def createLogic(attr: Attributes): GraphStageLogic = new GraphStageLogic(shape) { | ||
| var buffer = ByteString.empty | ||
|
|
||
| setHandler(out, new OutHandler { | ||
| override def onPull(): Unit = { | ||
| if (isClosed(in)) { | ||
| emit() | ||
| } else { | ||
| pull(in) | ||
| } | ||
| } | ||
| }) | ||
|
|
||
| setHandler(in, new InHandler { | ||
| override def onPush(): Unit = { | ||
| val elem = grab(in) | ||
| buffer ++= elem | ||
| pull(in) | ||
| } | ||
|
|
||
| override def onUpstreamFinish(): Unit = { | ||
| if (isAvailable(out)) { | ||
| emit() | ||
| } | ||
| completeStage() | ||
| } | ||
| }) | ||
|
|
||
| def emit(): Unit = { | ||
| emit(out, Chunk(Source.single(buffer), buffer.size), () => completeStage()) | ||
| } | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.