Apache Spark Network Common

import org.apache.spark.network.TransportContext;
import org.apache.spark.network.client.TransportClient;
import org.apache.spark.network.client.TransportClientFactory;
import org.apache.spark.network.server.TransportServer;
import org.apache.spark.network.server.RpcHandler;
import org.apache.spark.network.buffer.ManagedBuffer;
import org.apache.spark.network.util.TransportConf;

import org.apache.spark.network.TransportContext;
import org.apache.spark.network.server.RpcHandler;
import org.apache.spark.network.server.NoOpRpcHandler;
import org.apache.spark.network.util.TransportConf;
import org.apache.spark.network.util.SystemPropertyConfigProvider;

// Create transport configuration
TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());

// Create RPC handler
RpcHandler rpcHandler = new NoOpRpcHandler();

// Create transport context
TransportContext context = new TransportContext(conf, rpcHandler);

// Create client factory for outbound connections
TransportClientFactory clientFactory = context.createClientFactory();

// Create client to connect to remote server
TransportClient client = clientFactory.createClient("localhost", 8080);

// Create server for inbound connections
TransportServer server = context.createServer(8080, new ArrayList<>());

public class TransportContext {
  public TransportContext(TransportConf conf, RpcHandler rpcHandler);
  public TransportContext(TransportConf conf, RpcHandler rpcHandler, boolean closeIdleConnections);
  
  public TransportClientFactory createClientFactory();
  public TransportClientFactory createClientFactory(List<TransportClientBootstrap> bootstraps);
  
  public TransportServer createServer();
  public TransportServer createServer(int port, List<TransportServerBootstrap> bootstraps);
  public TransportServer createServer(String host, int port, List<TransportServerBootstrap> bootstraps);
}

public class TransportClient implements Closeable {
  public void fetchChunk(long streamId, int chunkIndex, ChunkReceivedCallback callback);
  public void stream(String streamId, StreamCallback callback);
  public long sendRpc(ByteBuffer message, RpcResponseCallback callback);
  public ByteBuffer sendRpcSync(ByteBuffer message, long timeoutMs);
  public void send(ByteBuffer message);
}

public class TransportClientFactory implements Closeable {
  public TransportClient createClient(String remoteHost, int remotePort);
  public TransportClient createUnmanagedClient(String remoteHost, int remotePort);  
}

public class TransportServer implements Closeable {
  public int getPort();
  public void close();
}

public abstract class RpcHandler {
  public abstract void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback);
  public abstract StreamManager getStreamManager();
}

public abstract class StreamManager {
  public abstract ManagedBuffer getChunk(long streamId, int chunkIndex);
}

public interface Message {
  Type type();
  ManagedBuffer body();
  boolean isBodyInFrame();
  
  enum Type {
    ChunkFetchRequest, ChunkFetchSuccess, ChunkFetchFailure,
    RpcRequest, RpcResponse, RpcFailure, 
    StreamRequest, StreamResponse, StreamFailure,
    OneWayMessage, User
  }
}

public class RpcRequest extends AbstractMessage implements RequestMessage {
  public final long requestId;
  public RpcRequest(long requestId, ManagedBuffer message);
}

public class ChunkFetchRequest extends AbstractMessage implements RequestMessage {
  public final StreamChunkId streamChunkId;
  public ChunkFetchRequest(StreamChunkId streamChunkId);
}

public abstract class ManagedBuffer {
  public abstract long size();
  public abstract ByteBuffer nioByteBuffer() throws IOException;
  public abstract InputStream createInputStream() throws IOException;
  public abstract ManagedBuffer retain();
  public abstract ManagedBuffer release();
  public abstract Object convertToNetty() throws IOException;
}

public class NioManagedBuffer extends ManagedBuffer {
  public NioManagedBuffer(ByteBuffer buf);
}

public class FileSegmentManagedBuffer extends ManagedBuffer {
  public FileSegmentManagedBuffer(TransportConf conf, File file, long offset, long length);
}

public interface SecretKeyHolder {
  String getSaslUser(String appId);
  String getSecretKey(String appId);
}

public class SaslClientBootstrap implements TransportClientBootstrap {
  public SaslClientBootstrap(TransportConf conf, String appId, SecretKeyHolder secretKeyHolder);
}

public class SaslServerBootstrap implements TransportServerBootstrap {
  public SaslServerBootstrap(TransportConf conf, SecretKeyHolder secretKeyHolder);
}

public class TransportConf {
  public TransportConf(String module, ConfigProvider conf);
  
  public String ioMode();
  public boolean preferDirectBufs();
  public int connectionTimeoutMs();
  public int numConnectionsPerPeer(); 
  public int serverThreads();
  public int clientThreads();
}

public abstract class ConfigProvider {
  public abstract String get(String name);
  public String get(String name, String defaultValue);
  public int getInt(String name, int defaultValue);
  public boolean getBoolean(String name, boolean defaultValue);
}

public interface RequestMessage extends Message {}
public interface ResponseMessage extends Message {}

public interface ChunkReceivedCallback {
  void onSuccess(int chunkIndex, ManagedBuffer buffer);
  void onFailure(int chunkIndex, Throwable e);
}

public interface RpcResponseCallback {
  void onSuccess(ByteBuffer response);
  void onFailure(Throwable e);
}

public interface StreamCallback {
  void onData(String streamId, ByteBuffer buf);
  void onComplete(String streamId);
  void onFailure(String streamId, Throwable cause);
}

public interface TransportClientBootstrap {
  void doBootstrap(TransportClient client, Channel channel);
}

public interface TransportServerBootstrap {
  RpcHandler doBootstrap(Channel channel, RpcHandler rpcHandler);
}

public class StreamChunkId {
  public final long streamId;
  public final int chunkIndex;
  
  public StreamChunkId(long streamId, int chunkIndex);
  public String toString();
  public boolean equals(Object other);
  public int hashCode();
}

public enum IOMode {
  NIO, EPOLL
}

public enum ByteUnit {
  BYTE, KiB, MiB, GiB, TiB, PiB;
  
  public long convertFrom(long d, ByteUnit u);
  public long convertTo(long d, ByteUnit u);
  public double toBytes(long d);
  public long toKiB(long d);
  public long toMiB(long d);
  public long toGiB(long d);
  public long toTiB(long d);
  public long toPiB(long d);
}

public class ChunkFetchFailureException extends RuntimeException {
  public ChunkFetchFailureException(String errorMsg, Throwable cause);
  public ChunkFetchFailureException(String errorMsg);
}