private static void CreateAndVerifyConnection(ISocketWrapper daemonSocket)
{
var ipEndpoint = (IPEndPoint)daemonSocket.LocalEndPoint;
int port = ipEndpoint.Port;
ISocketWrapper clientSocket = SocketFactory.CreateSocket();
clientSocket.Connect(ipEndpoint.Address, port);
// Now process the bytes flowing in from the client.
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
payloadWriter.WriteTestData(clientSocket.OutputStream);
List <object[]> rowsReceived = PayloadReader.Read(clientSocket.InputStream);
// Validate rows received.
Assert.Equal(10, rowsReceived.Count);
for (int i = 0; i < 10; ++i)
{
// Two UDFs registered, thus expecting two columns.
// Refer to TestData.GetDefaultCommandPayload().
object[] row = rowsReceived[i];
Assert.Equal(2, rowsReceived[i].Length);
Assert.Equal($"udf2 udf1 {i}", row[0]);
Assert.Equal(i + i, row[1]);
}
}
public void TestsSimpleWorkerTaskRunners(string version)
{
using ISocketWrapper serverListener = SocketFactory.CreateSocket();
var ipEndpoint = (IPEndPoint)serverListener.LocalEndPoint;
serverListener.Listen();
var typedVersion = new Version(version);
var simpleWorker = new SimpleWorker(typedVersion);
Task clientTask = Task.Run(() => simpleWorker.Run(ipEndpoint.Port));
PayloadWriter payloadWriter = new PayloadWriterFactory().Create(typedVersion);
using (ISocketWrapper serverSocket = serverListener.Accept())
{
if ((typedVersion.Major == 3 && typedVersion.Minor >= 2) || typedVersion.Major > 3)
{
int pid = SerDe.ReadInt32(serverSocket.InputStream);
}
TaskRunnerTests.TestTaskRunnerReadWrite(serverSocket, payloadWriter);
}
Assert.True(clientTask.Wait(5000));
}
public void TestClosedStreamWithSocket()
{
var commandPayload = new CommandPayload()
{
EvalType = UdfUtils.PythonEvalType.SQL_BATCHED_UDF,
Commands = new Command[] { }
};
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
Payload payload = TestData.GetDefaultPayload();
var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
using (ISocketWrapper serverSocket = serverListener.Accept())
{
Stream outStream = serverSocket.OutputStream;
payloadWriter.Write(outStream, payload, commandPayload);
outStream.Flush();
}
// At this point server socket is closed.
Stream inStream = clientSocket.InputStream;
// Consume bytes already written to the socket.
var payloadProcessor = new PayloadProcessor(payloadWriter.Version);
Payload actualPayload = payloadProcessor.Process(inStream);
Assert.Equal(payload.SplitIndex, actualPayload.SplitIndex);
Assert.Equal(payload.Version, actualPayload.Version);
Assert.Equal(payload.TaskContext, actualPayload.TaskContext);
Assert.Equal(payload.SparkFilesDir, actualPayload.SparkFilesDir);
Assert.Equal(payload.IncludeItems, actualPayload.IncludeItems);
Assert.Equal(payload.BroadcastVariables.Count, actualPayload.BroadcastVariables.Count);
ValidateCommandPayload(commandPayload, actualPayload.Command);
// Another read will detect that the socket is closed.
Assert.Null(payloadProcessor.Process(inStream));
}
public void TestPayloadProcessor(string version)
{
CommandPayload commandPayload = TestData.GetDefaultCommandPayload();
PayloadWriter payloadWriter = new PayloadWriterFactory().Create(new Version(version));
Payload payload = TestData.GetDefaultPayload();
Payload actualPayload = null;
using (var outStream = new MemoryStream())
{
payloadWriter.Write(outStream, payload, commandPayload);
using (var inputStream = new MemoryStream(outStream.ToArray()))
{
actualPayload =
new PayloadProcessor(payloadWriter.Version).Process(inputStream);
}
}
// Validate the read payload.
Assert.Equal(payload.SplitIndex, actualPayload.SplitIndex);
Assert.Equal(payload.Version, actualPayload.Version);
Assert.Equal(payload.TaskContext, actualPayload.TaskContext);
Assert.Equal(payload.SparkFilesDir, actualPayload.SparkFilesDir);
Assert.Equal(payload.IncludeItems, actualPayload.IncludeItems);
Assert.Equal(payload.BroadcastVariables.Count, actualPayload.BroadcastVariables.Count);
ValidateCommandPayload(commandPayload, actualPayload.Command);
// Validate the UDFs.
var actualCommand1 = (SqlCommand)actualPayload.Command.Commands[0];
var result1 = ((PicklingWorkerFunction)actualCommand1.WorkerFunction).Func(
0,
new object[] { "hello", 10, 20 },
actualCommand1.ArgOffsets);
Assert.Equal("udf2 udf1 hello", result1);
var actualCommand2 = (SqlCommand)actualPayload.Command.Commands[1];
var result2 = ((PicklingWorkerFunction)actualCommand2.WorkerFunction).Func(
0,
new object[] { "hello", 10, 20 },
actualCommand2.ArgOffsets);
Assert.Equal(30, result2);
}
public void TestTaskRunner()
{
using var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
var taskRunner = new TaskRunner(0, clientSocket, false, payloadWriter.Version);
Task clientTask = Task.Run(() => taskRunner.Run());
TestTaskRunnerReadWrite(serverListener, payloadWriter);
Assert.True(clientTask.Wait(5000));
}
public void TestsSimpleWorkerTaskRunners(string version)
{
using ISocketWrapper serverListener = SocketFactory.CreateSocket();
var ipEndpoint = (IPEndPoint)serverListener.LocalEndPoint;
serverListener.Listen();
var typedVersion = new Version(version);
var simpleWorker = new SimpleWorker(typedVersion);
Task clientTask = Task.Run(() => simpleWorker.Run(ipEndpoint.Port));
PayloadWriter payloadWriter = new PayloadWriterFactory().Create(typedVersion);
TaskRunnerTests.TestTaskRunnerReadWrite(serverListener, payloadWriter);
Assert.True(clientTask.Wait(5000));
}
public void TestTaskRunner()
{
using var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
var taskRunner = new TaskRunner(0, clientSocket, false, payloadWriter.Version);
Task clientTask = Task.Run(() => taskRunner.Run());
using (ISocketWrapper serverSocket = serverListener.Accept())
{
System.IO.Stream inputStream = serverSocket.InputStream;
System.IO.Stream outputStream = serverSocket.OutputStream;
payloadWriter.WriteTestData(outputStream);
// Now process the bytes flowing in from the client.
List <object[]> rowsReceived = PayloadReader.Read(inputStream);
// Validate rows received.
Assert.Equal(10, rowsReceived.Count);
for (int i = 0; i < 10; ++i)
{
// Two UDFs registered, thus expecting two columns.
// Refer to TestData.GetDefaultCommandPayload().
object[] row = rowsReceived[i];
Assert.Equal(2, rowsReceived[i].Length);
Assert.Equal($"udf2 udf1 {i}", row[0]);
Assert.Equal(i + i, row[1]);
}
}
Assert.True(clientTask.Wait(5000));
}
public void TestTaskRunner()
{
using var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
var taskRunner = new TaskRunner(0, clientSocket, false, payloadWriter.Version);
var clientTask = Task.Run(() => taskRunner.Run());
using (ISocketWrapper serverSocket = serverListener.Accept())
{
System.IO.Stream inputStream = serverSocket.InputStream;
System.IO.Stream outputStream = serverSocket.OutputStream;
Payload payload = TestData.GetDefaultPayload();
CommandPayload commandPayload = TestData.GetDefaultCommandPayload();
payloadWriter.Write(outputStream, payload, commandPayload);
// Write 10 rows to the output stream.
var pickler = new Pickler();
for (int i = 0; i < 10; ++i)
{
var pickled = pickler.dumps(
new[] { new object[] { i.ToString(), i, i } });
SerDe.Write(outputStream, pickled.Length);
SerDe.Write(outputStream, pickled);
}
// Signal the end of data and stream.
SerDe.Write(outputStream, (int)SpecialLengths.END_OF_DATA_SECTION);
SerDe.Write(outputStream, (int)SpecialLengths.END_OF_STREAM);
outputStream.Flush();
// Now process the bytes flowing in from the client.
var timingDataReceived = false;
var exceptionThrown = false;
var rowsReceived = new List <object[]>();
while (true)
{
var length = SerDe.ReadInt32(inputStream);
if (length > 0)
{
var pickledBytes = SerDe.ReadBytes(inputStream, length);
using var unpickler = new Unpickler();
var rows = unpickler.loads(pickledBytes) as ArrayList;
foreach (object row in rows)
{
rowsReceived.Add((object[])row);
}
}
else if (length == (int)SpecialLengths.TIMING_DATA)
{
var bootTime = SerDe.ReadInt64(inputStream);
var initTime = SerDe.ReadInt64(inputStream);
var finishTime = SerDe.ReadInt64(inputStream);
var memoryBytesSpilled = SerDe.ReadInt64(inputStream);
var diskBytesSpilled = SerDe.ReadInt64(inputStream);
timingDataReceived = true;
}
else if (length == (int)SpecialLengths.PYTHON_EXCEPTION_THROWN)
{
SerDe.ReadString(inputStream);
exceptionThrown = true;
break;
}
else if (length == (int)SpecialLengths.END_OF_DATA_SECTION)
{
var numAccumulatorUpdates = SerDe.ReadInt32(inputStream);
SerDe.ReadInt32(inputStream);
break;
}
}
Assert.True(timingDataReceived);
Assert.False(exceptionThrown);
// Validate rows received.
Assert.Equal(10, rowsReceived.Count);
for (int i = 0; i < 10; ++i)
{
// Two UDFs registered, thus expecting two columns.
// Refer to TestData.GetDefaultCommandPayload().
var row = rowsReceived[i];
Assert.Equal(2, rowsReceived[i].Length);
Assert.Equal($"udf2 udf1 {i}", row[0]);
Assert.Equal(i + i, row[1]);
}
}
Assert.True(clientTask.Wait(5000));
}
