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 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 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));
}
