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