From be895232e1c7b203f52a9d467c0e81236253bdfa Mon Sep 17 00:00:00 2001 From: Stefan Zeiger Date: Fri, 21 Aug 2020 17:58:21 +0200 Subject: [PATCH] Write output directly to disk with -o This prevents caching the entire output string in memory, reducing the heap usage for our largest output file of 7.2 MB from 101 MB to 93 MB in my test. Only JSON is supported for now. YAML requires more work because the YamlRenderer is more closely tied to the underlying StringWriter, so it still has to buffer. --- sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala | 47 +++++++++++-------- sjsonnet/src/sjsonnet/Renderer.scala | 14 +++--- .../test/src-jvm/sjsonnet/MainTests.scala | 32 +++++++++++++ 3 files changed, 67 insertions(+), 26 deletions(-) create mode 100644 sjsonnet/test/src-jvm/sjsonnet/MainTests.scala diff --git a/sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala b/sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala index 4cec3e88..c52ab15f 100644 --- a/sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala +++ b/sjsonnet/src-jvm/sjsonnet/SjsonnetMain.scala @@ -1,6 +1,7 @@ package sjsonnet -import java.io.{BufferedOutputStream, InputStream, OutputStreamWriter, PrintStream, StringWriter} +import java.io.{BufferedOutputStream, InputStream, OutputStreamWriter, PrintStream, StringWriter, Writer} +import java.nio.charset.StandardCharsets import java.nio.file.NoSuchFileException import sjsonnet.Cli.Config @@ -75,10 +76,10 @@ object SjsonnetMain { result match{ case Left(err) => - if (!err.isEmpty) System.err.println(err) + if (!err.isEmpty) stderr.println(err) 1 case Right(str) => - if (!str.isEmpty) System.out.println(str) + if (!str.isEmpty) stdout.println(str) 0 } } @@ -105,15 +106,25 @@ object SjsonnetMain { config.preserveOrder ) - def writeFile(f: os.RelPath, contents: String): Either[String, Unit] = { - Try(os.write.over(os.Path(f, wd), contents, createFolders = config.createDirs)) - .toEither - .left - .map{ - case e: NoSuchFileException => s"open $f: no such file or directory" - case e => e.toString - } - } + def handleWriteFile[T](f: => T): Either[String, T] = + Try(f).toEither.left.map{ + case e: NoSuchFileException => s"open $f: no such file or directory" + case e => e.toString + } + + def writeFile(f: os.RelPath, contents: String): Either[String, Unit] = + handleWriteFile(os.write.over(os.Path(f, wd), contents, createFolders = config.createDirs)) + + def writeToFile[U](f: os.RelPath, contents: Writer => Either[String, U]): Either[String, Unit] = + handleWriteFile(os.write.over.outputStream(os.Path(f, wd), createFolders = config.createDirs)).flatMap { out => + try { + val buf = new BufferedOutputStream(out) + val wr = new OutputStreamWriter(buf, StandardCharsets.UTF_8) + val u = contents(wr) + wr.flush() + u.map(_ => ()) + } finally out.close() + } (config.multi, config.yamlStream) match { case (Some(multiPath), _) => @@ -167,22 +178,20 @@ object SjsonnetMain { "whose elements hold the JSON for each document in the stream.") } case _ => - val materialized = interp.interpret0( + def materialize(wr: Writer) = interp.interpret0( os.read(path), OsPath(path), - new Renderer(indent = config.indent) + new Renderer(wr, indent = config.indent) ) config.outputFile match{ - case None => materialized.map(_.toString) + case None => + materialize(new StringWriter).map(_.toString) case Some(f) => val filePath = os.FilePath(f) match{ case _: os.Path => os.Path(f).relativeTo(os.pwd) case _ => os.RelPath(f) } - for{ - materializedStr <- materialized - _ <- writeFile(filePath, materializedStr.toString) - } yield "" + writeToFile(filePath, materialize(_)).map(_ => "") } } } diff --git a/sjsonnet/src/sjsonnet/Renderer.scala b/sjsonnet/src/sjsonnet/Renderer.scala index de4946c8..b794c1bc 100644 --- a/sjsonnet/src/sjsonnet/Renderer.scala +++ b/sjsonnet/src/sjsonnet/Renderer.scala @@ -1,5 +1,5 @@ package sjsonnet -import java.io.StringWriter +import java.io.{StringWriter, Writer} import java.util.regex.Pattern import upickle.core.{ArrVisitor, ObjVisitor} @@ -13,7 +13,7 @@ import ujson.BaseRenderer * - Custom printing of empty dictionaries and arrays * */ -class Renderer(out: StringWriter = new java.io.StringWriter(), +class Renderer(out: Writer = new java.io.StringWriter(), indent: Int = -1) extends BaseRenderer(out, indent){ var newlineBuffered = false override def visitFloat64(d: Double, index: Int) = { @@ -40,7 +40,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(), newlineBuffered = false commaBuffered = false } - override def visitArray(length: Int, index: Int) = new ArrVisitor[StringWriter, StringWriter] { + override def visitArray(length: Int, index: Int) = new ArrVisitor[Writer, Writer] { var empty = true flushBuffer() out.append('[') @@ -48,7 +48,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(), depth += 1 def subVisitor = Renderer.this - def visitValue(v: StringWriter, index: Int): Unit = { + def visitValue(v: Writer, index: Int): Unit = { empty = false flushBuffer() commaBuffered = true @@ -65,7 +65,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(), } } - override def visitObject(length: Int, index: Int) = new ObjVisitor[StringWriter, StringWriter] { + override def visitObject(length: Int, index: Int) = new ObjVisitor[Writer, Writer] { var empty = true flushBuffer() out.append('{') @@ -78,7 +78,7 @@ class Renderer(out: StringWriter = new java.io.StringWriter(), flushBuffer() out.append(colonSnippet) } - def visitValue(v: StringWriter, index: Int): Unit = { + def visitValue(v: Writer, index: Int): Unit = { commaBuffered = true } def visitEnd(index: Int) = { @@ -212,7 +212,7 @@ class YamlRenderer(out: StringWriter = new java.io.StringWriter(), indentArrayIn } } -class PythonRenderer(out: StringWriter = new java.io.StringWriter(), +class PythonRenderer(out: Writer = new java.io.StringWriter(), indent: Int = -1) extends BaseRenderer(out, indent){ override def visitNull(index: Int) = { diff --git a/sjsonnet/test/src-jvm/sjsonnet/MainTests.scala b/sjsonnet/test/src-jvm/sjsonnet/MainTests.scala new file mode 100644 index 00000000..ffdf3ec3 --- /dev/null +++ b/sjsonnet/test/src-jvm/sjsonnet/MainTests.scala @@ -0,0 +1,32 @@ +package sjsonnet + +import java.io.{ByteArrayOutputStream, File, PrintStream} +import java.util.Arrays + +import utest._ + +object MainTests extends TestSuite { + + val testSuiteRoot = os.pwd / "sjsonnet" / "test" / "resources" / "test_suite" + + val tests = Tests { + // Compare writing to stdout with writing to a file + test("writeToFile") { + val source = (testSuiteRoot / "local.jsonnet").toString() + val outF = File.createTempFile("sjsonnet", ".json") + val out = new ByteArrayOutputStream() + val pout = new PrintStream(out) + SjsonnetMain.main0(Array(source), collection.mutable.Map.empty, System.in, pout, System.err, os.pwd, None) + pout.flush() + SjsonnetMain.main0(Array("-o", outF.getAbsolutePath, source), collection.mutable.Map.empty, System.in, System.out, System.err, os.pwd, None) + val stdoutBytes = out.toByteArray + val fileBytes = os.read(os.Path(outF)).getBytes + // stdout mode uses println so it has an extra platform-specific line separator at the end + val eol = System.getProperty("line.separator").getBytes + + //println(stdoutBytes.map(_.toInt).mkString(",")) + //println(fileBytes.map(_.toInt).mkString(",")) + assert(Arrays.equals(fileBytes ++ eol, stdoutBytes)) + } + } +}