Add @context support for JStachio module

modules/jooby-jstachio/src/main/java/io/jooby/jstachio/JStachioBuffer.java

@@ -5,16 +5,20 @@
  */
 package io.jooby.jstachio;
 
+import java.nio.charset.StandardCharsets;
+
+import io.jstach.jstachio.output.ByteBufferEncodedOutput;
+
 /**
  * To provide Rocker like buffer support
  *
  * @author agentgt
  */
 interface JStachioBuffer {
 
-  public ByteBufferedOutputStream acquire();
+  public ByteBufferEncodedOutput acquire();
 
-  public void release(ByteBufferedOutputStream buffer);
+  public void release(ByteBufferEncodedOutput buffer);
 
   static JStachioBuffer of(int bufferSize, boolean reuseBuffer) {
     if (reuseBuffer) {
@@ -27,29 +31,31 @@ static JStachioBuffer of(int bufferSize, boolean reuseBuffer) {
 
 record NoReuseBuffer(int bufferSize) implements JStachioBuffer {
   @Override
-  public ByteBufferedOutputStream acquire() {
-    return new ByteBufferedOutputStream(bufferSize);
+  public ByteBufferEncodedOutput acquire() {
+    return ByteBufferEncodedOutput.ofByteArray(StandardCharsets.UTF_8, bufferSize);
   }
 
   @Override
-  public void release(ByteBufferedOutputStream buffer) {}
+  public void release(ByteBufferEncodedOutput buffer) {}
 }
 
 class ReuseBuffer implements JStachioBuffer {
-  private final ThreadLocal<ByteBufferedOutputStream> localBuffer;
+  private final ThreadLocal<ByteBufferEncodedOutput> localBuffer;
 
   public ReuseBuffer(int bufferSize) {
     super();
-    this.localBuffer = ThreadLocal.withInitial(() -> new ByteBufferedOutputStream(bufferSize));
+    this.localBuffer =
+        ThreadLocal.withInitial(
+            () -> ByteBufferEncodedOutput.ofByteArray(StandardCharsets.UTF_8, bufferSize));
   }
 
   @Override
-  public ByteBufferedOutputStream acquire() {
+  public ByteBufferEncodedOutput acquire() {
     return localBuffer.get();
   }
 
   @Override
-  public void release(ByteBufferedOutputStream buffer) {
-    buffer.reset();
+  public void release(ByteBufferEncodedOutput buffer) {
+    buffer.close();
   }
 }

modules/jooby-jstachio/src/main/java/io/jooby/jstachio/JStachioHandler.java

@@ -6,18 +6,21 @@
 package io.jooby.jstachio;
 
 import java.io.IOException;
+import java.util.function.BiFunction;
 
 import io.jooby.Context;
-import io.jooby.MediaType;
 import io.jooby.Route;
 import io.jooby.Route.Handler;
 import io.jstach.jstachio.JStachio;
-import io.jstach.jstachio.Template;
+import io.jstach.jstachio.output.ByteBufferEncodedOutput;
 
 class JStachioHandler extends JStachioRenderer<Context> implements Route.Filter {
 
-  public JStachioHandler(JStachio jstachio, JStachioBuffer buffer) {
-    super(jstachio, buffer);
+  public JStachioHandler(
+      JStachio jstachio,
+      JStachioBuffer buffer,
+      BiFunction<Context, String, String> contextFunction) {
+    super(jstachio, buffer, contextFunction);
   }
 
   @Override
@@ -33,24 +36,9 @@ public Handler apply(Handler next) {
     };
   }
 
-  @SuppressWarnings("unchecked")
   @Override
-  Context render(
-      Context ctx,
-      @SuppressWarnings("rawtypes") Template template,
-      Object model,
-      ByteBufferedOutputStream stream)
-      throws IOException {
-    ctx.setResponseType(MediaType.html);
-    template.write(model, stream);
-    /*
-     * Rocker used a byte buffer here BUT it just wraps the internal buffer in the stream
-     * instead of copying.
-     *
-     * Which is good for performance but bad if the ctx.send call is not blocking aka
-     * hand the buffer off to another thread.
-     */
-    ctx.send(stream.toBuffer());
+  Context extractOutput(Context ctx, ByteBufferEncodedOutput stream) throws IOException {
+    ctx.send(stream.asByteBuffer());
     return ctx;
   }
 }

modules/jooby-jstachio/src/main/java/io/jooby/jstachio/JStachioMessageEncoder.java

@@ -6,19 +6,20 @@
 package io.jooby.jstachio;
 
 import java.io.IOException;
+import java.util.function.BiFunction;
 
 import io.jooby.Context;
 import io.jooby.MessageEncoder;
 import io.jstach.jstachio.JStachio;
-import io.jstach.jstachio.Template;
+import io.jstach.jstachio.output.ByteBufferEncodedOutput;
 
-/*
- * Should this be public?
- */
 class JStachioMessageEncoder extends JStachioRenderer<byte[]> implements MessageEncoder {
 
-  public JStachioMessageEncoder(JStachio jstachio, JStachioBuffer buffer) {
-    super(jstachio, buffer);
+  public JStachioMessageEncoder(
+      JStachio jstachio,
+      JStachioBuffer buffer,
+      BiFunction<Context, String, String> contextFunction) {
+    super(jstachio, buffer, contextFunction);
   }
 
   @Override
@@ -29,15 +30,8 @@ public byte[] encode(Context ctx, Object value) throws Exception {
     return null;
   }
 
-  @SuppressWarnings("unchecked")
   @Override
-  byte[] render(
-      Context ctx,
-      @SuppressWarnings("rawtypes") Template template,
-      Object model,
-      ByteBufferedOutputStream stream)
-      throws IOException {
-    template.write(model, stream);
+  byte[] extractOutput(Context ctx, ByteBufferEncodedOutput stream) throws IOException {
     return stream.toByteArray();
   }
 }

modules/jooby-jstachio/src/main/java/io/jooby/jstachio/JStachioModule.java

@@ -6,9 +6,11 @@
 package io.jooby.jstachio;
 
 import java.util.ServiceLoader;
+import java.util.function.BiFunction;
 
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
+import io.jooby.Context;
 import io.jooby.Extension;
 import io.jooby.Jooby;
 import io.jooby.ServiceRegistry;
@@ -34,6 +36,7 @@ public class JStachioModule implements Extension {
   private @Nullable JStachio jstachio;
   private int bufferSize = 8 * 1024;
   private boolean reuseBuffer;
+  private BiFunction<Context, String, String> contextFunction = (ctx, key) -> ctx.getAttribute(key);
 
   /**
    * Sets the jstachio to use instead of the default.
@@ -73,6 +76,21 @@ public JStachioModule reuseBuffer(boolean reuseBuffer) {
     return this;
   }
 
+  /**
+   * JStachio will by default bind {@linkplain Context#getAttributes() Context attributes} to <code>
+   * &#64;context</code>. This configuration option allows fetching context keys from something
+   * else. <a
+   * href="https://jstach.io/jstachio/io.jstach.jstachio/io/jstach/jstachio/context/package-summary.html">
+   * See JStachio doc on context. </a>
+   *
+   * @param contextFunction
+   * @return This module.
+   */
+  public JStachioModule contextFunction(BiFunction<Context, String, String> contextFunction) {
+    this.contextFunction = contextFunction;
+    return this;
+  }
+
   /**
    * Installs JStachio into Jooby and provides the JStachio service to the {@link ServiceRegistry}.
    * {@inheritDoc}
@@ -98,8 +116,8 @@ public void install(@NonNull Jooby application) throws Exception {
     }
     JStachioBuffer buffer = JStachioBuffer.of(bufferSize, reuseBuffer);
 
-    JStachioMessageEncoder encoder = new JStachioMessageEncoder(j, buffer);
-    JStachioResultHandler handler = new JStachioResultHandler(j, buffer);
+    JStachioMessageEncoder encoder = new JStachioMessageEncoder(j, buffer, contextFunction);
+    JStachioResultHandler handler = new JStachioResultHandler(j, buffer, contextFunction);
     application.encoder(encoder);
     application.resultHandler(handler);
   }

modules/jooby-jstachio/src/main/java/io/jooby/jstachio/JStachioRenderer.java

@@ -6,11 +6,14 @@
 package io.jooby.jstachio;
 
 import java.io.IOException;
+import java.util.function.BiFunction;
 
 import io.jooby.Context;
 import io.jooby.MediaType;
 import io.jstach.jstachio.JStachio;
-import io.jstach.jstachio.Template;
+import io.jstach.jstachio.context.ContextJStachio;
+import io.jstach.jstachio.context.ContextNode;
+import io.jstach.jstachio.output.ByteBufferEncodedOutput;
 
 /**
  * Shared logic between encoder and result handler
@@ -20,20 +23,23 @@
  */
 abstract class JStachioRenderer<T> {
 
-  private final JStachio jstachio;
+  private final ContextJStachio jstachio;
   private final JStachioBuffer buffer;
+  private final BiFunction<Context, String, String> contextFunction;
 
-  public JStachioRenderer(JStachio jstachio, JStachioBuffer buffer) {
+  public JStachioRenderer(
+      JStachio jstachio,
+      JStachioBuffer buffer,
+      BiFunction<Context, String, String> contextFunction) {
     super();
-    this.jstachio = jstachio;
+    this.jstachio = ContextJStachio.of(jstachio);
     this.buffer = buffer;
+    this.contextFunction = contextFunction;
   }
 
   public T render(Context ctx, Object model) throws Exception {
     var stream = buffer.acquire();
     try {
-      @SuppressWarnings("rawtypes")
-      Template template = jstachio.findTemplate(model);
       /*
        * TODO we probably should resolve the correct media type here and more importantly charset
        * Or at least validate that it is text/html UTF-8.
@@ -43,18 +49,15 @@ public T render(Context ctx, Object model) throws Exception {
        * So for now we will as well.
        */
       ctx.setResponseType(MediaType.html);
-      return render(ctx, template, model, stream);
+      ContextNode contextNode = ContextNode.of(s -> contextFunction.apply(ctx, s));
+      jstachio.write(model, contextNode, stream);
+      return extractOutput(ctx, stream);
     } finally {
       buffer.release(stream);
     }
   }
 
-  abstract T render(
-      Context ctx,
-      @SuppressWarnings("rawtypes") Template template,
-      Object model,
-      ByteBufferedOutputStream stream)
-      throws IOException;
+  abstract T extractOutput(Context ctx, ByteBufferEncodedOutput stream) throws IOException;
 
   protected boolean supportsType(Class<?> modelClass) {
     return jstachio.supportsType(modelClass);