Skip to content

Commit

Permalink
Enable json streaming mode in AgentServlet. Fixes #282
Browse files Browse the repository at this point in the history
  • Loading branch information
@paoloantinori
paoloantinori committed Sep 8, 2016
1 parent b6010e3 commit ae1bc77
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 27 deletions.
86 changes: 76 additions & 10 deletions agent/core/src/main/java/org/jolokia/http/AgentServlet.java
View file
Expand Up @@ -18,6 +18,8 @@
import org.jolokia.restrictor.*;
import org.jolokia.util.*;
import org.json.simple.JSONAware;
import org.json.simple.JSONStreamAware;


/*
* Copyright 2009-2013 Roland Huss
Expand Down Expand Up @@ -77,6 +79,9 @@ public class AgentServlet extends HttpServlet {
// whether to allow reverse DNS lookup for checking the remote host
private boolean allowDnsReverseLookup;

// wheter to allow streaming mode for response
private boolean streamingEnabled = true;

/**
* No argument constructor, used e.g. by an servlet
* descriptor when creating the servlet out of web.xml
Expand Down Expand Up @@ -135,6 +140,7 @@ public void init(ServletConfig pServletConfig) throws ServletException {
backendManager = new BackendManager(config,logHandler, restrictor);
requestHandler = new HttpRequestHandler(config,backendManager,logHandler);
allowDnsReverseLookup = config.getAsBoolean(ConfigKey.ALLOW_DNS_REVERSE_LOOKUP);
streamingEnabled = config.getAsBoolean(ConfigKey.STREAMING);

initDiscoveryMulticast(config);
}
Expand Down Expand Up @@ -266,15 +272,19 @@ private void handle(ServletRequestHandler pReqHandler,HttpServletRequest pReq, H
setCorsHeader(pReq, pResp);

String callback = pReq.getParameter(ConfigKey.CALLBACK.getKeyValue());
String answer = json != null ?
json.toJSONString() :
requestHandler.handleThrowable(new Exception("Internal error while handling an exception")).toJSONString();
json = json != null ?
json :
requestHandler.handleThrowable(new Exception("Internal error while handling an exception"));

String mimeType;
if (callback != null) {
// Send a JSONP response
sendResponse(pResp, "text/javascript", callback + "(" + answer + ");");
mimeType = "text/javascript";
} else {
sendResponse(pResp, getMimeType(pReq),answer);
mimeType = getMimeType(pReq);
}

sendResponse(pResp, mimeType, callback, json);
}
}

Expand Down Expand Up @@ -449,12 +459,68 @@ Configuration initConfig(ServletConfig pConfig) {
return config;
}

private void sendResponse(HttpServletResponse pResp, String pContentType, String pJsonTxt) throws IOException {
setContentType(pResp, pContentType);
pResp.setStatus(200);
private void sendResponse(HttpServletResponse pResp, String mimeType, String callback, JSONAware pJson) throws IOException {
setContentType(pResp, mimeType);
pResp.setStatus(HttpServletResponse.SC_OK);
setNoCacheHeaders(pResp);
PrintWriter writer = pResp.getWriter();
writer.write(pJsonTxt);
if (streamingEnabled) {
JSONStreamAware jsonStream = (JSONStreamAware)pJson;
sendStreamingResponse(pResp, callback, jsonStream);
} else {
// Fallback, send as one object
// TODO: Remove for 2.0
sendAllJSON(pResp, callback, pJson);
}
}

private void sendStreamingResponse(HttpServletResponse pResp, String callback, JSONStreamAware pJson) throws IOException {
ChunkedWriter writer = null;
try {
if (pJson != null) {
writer = new ChunkedWriter(pResp.getOutputStream(), "UTF-8");
if (callback == null) {
pJson.writeJSONString(writer);
} else {
writer.write(callback);
writer.write("(");
pJson.writeJSONString(writer);
writer.write(");");
}
} else {
pResp.setContentLength(-1);
}
} finally {
if (writer != null) {
// Always close in order to finish the request.
// Otherwise the thread blocks.
writer.flush();
writer.close();
}
}
}

private void sendAllJSON(HttpServletResponse pResp, String callback, JSONAware pJson) throws IOException {
OutputStream out = null;
try {
if (pJson != null) {
String json = pJson.toJSONString();
String content = callback == null ? json : callback + "(" + json + ");";
byte[] response = content.getBytes("UTF8");
pResp.setStatus(HttpServletResponse.SC_OK);
pResp.setContentLength(response.length);
out = pResp.getOutputStream();
out.write(response);
} else {
pResp.setStatus(HttpServletResponse.SC_OK);
pResp.setContentLength(-1);
}
} finally {
if (out != null) {
// Always close in order to finish the request.
// Otherwise the thread blocks.
out.close();
}
}
}

private void setNoCacheHeaders(HttpServletResponse pResp) {
Expand Down
46 changes: 29 additions & 17 deletions agent/core/src/test/java/org/jolokia/http/AgentServletTest.java
View file
Expand Up @@ -191,7 +191,7 @@ public void initWithAgentDiscoveryAndUrlCreationAfterGet() throws ServletExcepti
prepareStandardInitialisation(ConfigKey.DISCOVERY_ENABLED.getKeyValue(), "true");
try {
String url = "http://10.9.11.1:9876/jolokia";
StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
getDiscoveryRequestSetup(url),
getStandardResponseSetup());
replay(request, response);
Expand Down Expand Up @@ -239,7 +239,7 @@ public void error(String message, Throwable t) {
public void simpleGet() throws ServletException, IOException {
prepareStandardInitialisation();

StringWriter sw = initRequestResponseMocks();
ByteArrayOutputStream sw = initRequestResponseMocks();
expect(request.getPathInfo()).andReturn(HttpTestUtil.HEAP_MEMORY_GET_REQUEST);
expect(request.getParameter(ConfigKey.MIME_TYPE.getKeyValue())).andReturn("text/plain");
expect(request.getAttribute("subject")).andReturn(null);
Expand Down Expand Up @@ -267,7 +267,7 @@ private void checkNoReverseDns(boolean enabled, String ... expectedHosts) throws
ConfigKey.RESTRICTOR_CLASS.getKeyValue(),NoDnsLookupRestrictorChecker.class.getName(),
ConfigKey.ALLOW_DNS_REVERSE_LOOKUP.getKeyValue(),Boolean.toString(enabled));
NoDnsLookupRestrictorChecker.expectedHosts = expectedHosts;
StringWriter sw = initRequestResponseMocks();
ByteArrayOutputStream sw = initRequestResponseMocks();
expect(request.getPathInfo()).andReturn(HttpTestUtil.HEAP_MEMORY_GET_REQUEST);
expect(request.getParameter(ConfigKey.MIME_TYPE.getKeyValue())).andReturn("text/plain");
expect(request.getAttribute("subject")).andReturn(null);
Expand Down Expand Up @@ -308,7 +308,7 @@ public boolean isRemoteAccessAllowed(String... pHostOrAddress) {
@Test
public void simpleGetWithUnsupportedGetParameterMapCall() throws ServletException, IOException {
prepareStandardInitialisation();
StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
new Runnable() {
public void run() {
expect(request.getHeader("Origin")).andStubReturn(null);
Expand Down Expand Up @@ -340,7 +340,7 @@ public void run() {
public void simplePost() throws ServletException, IOException {
prepareStandardInitialisation();

StringWriter responseWriter = initRequestResponseMocks();
ByteArrayOutputStream responseWriter = initRequestResponseMocks();
expect(request.getCharacterEncoding()).andReturn("utf-8");
expect(request.getParameter(ConfigKey.MIME_TYPE.getKeyValue())).andReturn("text/plain");
expect(request.getAttribute("subject")).andReturn(null);
Expand All @@ -359,7 +359,7 @@ public void simplePost() throws ServletException, IOException {
public void unknownMethodWhenSettingContentType() throws ServletException, IOException {
prepareStandardInitialisation();

StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
getStandardRequestSetup(),
new Runnable() {
public void run() {
Expand Down Expand Up @@ -423,7 +423,7 @@ public void corsHeaderGetCheckWithNullOrigin() throws ServletException, IOExcept
private void checkCorsGetOrigin(final String in, final String out) throws ServletException, IOException {
prepareStandardInitialisation();

StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
new Runnable() {
public void run() {
expect(request.getHeader("Origin")).andStubReturn(in);
Expand Down Expand Up @@ -469,7 +469,7 @@ private void setNoCacheHeaders(HttpServletResponse pResp) {
public void withCallback() throws IOException, ServletException {
prepareStandardInitialisation();

StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
"myCallback",
getStandardRequestSetup(),
new Runnable() {
Expand All @@ -496,7 +496,7 @@ public void withException() throws ServletException, IOException {
initConfigMocks(null, null,"Error 500", IllegalStateException.class);
replay(config, context);
servlet.init(config);
StringWriter sw = initRequestResponseMocks(
ByteArrayOutputStream sw = initRequestResponseMocks(
new Runnable() {
public void run() {
expect(request.getHeader("Origin")).andReturn(null);
Expand Down Expand Up @@ -533,7 +533,7 @@ public void debug() throws IOException, ServletException {

servlet.init(config);

StringWriter sw = initRequestResponseMocks();
ByteArrayOutputStream sw = initRequestResponseMocks();
expect(request.getPathInfo()).andReturn(HttpTestUtil.HEAP_MEMORY_GET_REQUEST);
expect(request.getParameter(ConfigKey.MIME_TYPE.getKeyValue())).andReturn(null);
expect(request.getAttribute("subject")).andReturn(null);
Expand Down Expand Up @@ -587,17 +587,17 @@ private void initConfigMocks(String[] pInitParams, String[] pContextParams,Strin
expectLastCall().anyTimes();
}

private StringWriter initRequestResponseMocks() throws IOException {
private ByteArrayOutputStream initRequestResponseMocks() throws IOException {
return initRequestResponseMocks(
getStandardRequestSetup(),
getStandardResponseSetup());
}

private StringWriter initRequestResponseMocks(Runnable requestSetup,Runnable responseSetup) throws IOException {
private ByteArrayOutputStream initRequestResponseMocks(Runnable requestSetup,Runnable responseSetup) throws IOException {
return initRequestResponseMocks(null,requestSetup,responseSetup);
}

private StringWriter initRequestResponseMocks(String callback,Runnable requestSetup,Runnable responseSetup) throws IOException {
private ByteArrayOutputStream initRequestResponseMocks(String callback,Runnable requestSetup,Runnable responseSetup) throws IOException {
request = createMock(HttpServletRequest.class);
response = createMock(HttpServletResponse.class);
setNoCacheHeaders(response);
Expand All @@ -606,10 +606,22 @@ private StringWriter initRequestResponseMocks(String callback,Runnable requestSe
requestSetup.run();
responseSetup.run();

StringWriter sw = new StringWriter();
PrintWriter writer = new PrintWriter(sw);
expect(response.getWriter()).andReturn(writer);
return sw;
class MyServletOutputStream extends ServletOutputStream {
ByteArrayOutputStream baos;
public void write(int b) throws IOException {
baos.write(b);
}

public void setBaos(ByteArrayOutputStream baos){
this.baos = baos;
}
}
ByteArrayOutputStream baos = new ByteArrayOutputStream();
MyServletOutputStream sos = new MyServletOutputStream();
sos.setBaos(baos);
expect(response.getOutputStream()).andReturn(sos);

return baos;
}

private void preparePostRequest(String pReq) throws IOException {
Expand Down

0 comments on commit ae1bc77

Please sign in to comment.