package org.asciidoctor.diagram;

import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.HashMap;
import java.util.Map;
import java.util.ServiceLoader;

/**
 * An simple, single client command server that accepts HTTP messages as input.
 */
public class CommandServer {
    private final ServerSocket serverSocket;
    private CommandProcessor processor;

    public static void main(String[] args) throws IOException {
        ServerSocket serverSocket = new ServerSocket(0);
        System.out.println(serverSocket.getLocalPort());
        System.out.flush();

        CommandServer server = new CommandServer(serverSocket);
        server.processRequests();
        server.terminate();
    }

    public CommandServer(ServerSocket socket) {
        this(socket, loadGenerators());
    }

    public CommandServer(ServerSocket socket, Map<String, DiagramGenerator> generators) {
        this.serverSocket = socket;
        this.processor = new CommandProcessor(generators);
    }

    static Map<String, DiagramGenerator> loadGenerators()
    {
        Map<String, DiagramGenerator> generatorMap = new HashMap<String, DiagramGenerator>();

        ServiceLoader<DiagramGenerator> generatorLoader = ServiceLoader.load(DiagramGenerator.class);
        for (DiagramGenerator generator : generatorLoader) {
            generatorMap.put(generator.getName(), generator);
        }
        return generatorMap;
    }

    public void processRequests() throws IOException {
        while (!serverSocket.isClosed()) {
            Socket client = serverSocket.accept();

            RequestInput input = new HTTPInputStream(client.getInputStream());
            ResponseOutput output = new HTTPOutputStream(client.getOutputStream());

            Request request;
            while ((request = input.readRequest()) != null) {
                output.writeResponse(processor.processRequest(request));
                if ("close".equals(request.headers.getValue(HTTPHeader.CONNECTION))) {
                    break;
                }
            }

            client.close();
        }
    }

    public void terminate() throws IOException {
        serverSocket.close();
    }
}

@@ -1278,16 +1278,54 @@
 or.class
+, CommandServer.class.getClassLoader()
 );%0A     
@@ -1444,32 +1444,275 @@
 tor);%0A        %7D%0A
+        ServiceLoader%3CDiagramGenerator%3E generatorLoaderTCCL = ServiceLoader.load(DiagramGenerator.class);%0A        for (DiagramGenerator generator : generatorLoaderTCCL) %7B%0A            generatorMap.put(generator.getName(), generator);%0A        %7D%0A
         return g

Load service impls with TCCL and current classloader