Как я реализовывал алгоритм маскирования для протокола WebSocket на Wolfram Language и подключал Си-библиотеку

Всем привет! Эта статья — про очень простой алгоритм маскирования данных при пересылке по протоколу WebSocket. Но рассказать я хочу не про сам алгоритм, а про путь оптимизации, который я прошел, чтобы сделать его эффективным. Я уверен, что можно еще лучше и если так, то надеюсь уважаемые читатели мне подскажут! Приступим…

WebSocket

Это протокол общения, который работает поверх HTTP. В момент создания соединения клиент отправляет серверу «рукопожатие» примерно вот такого вида:

GET / HTTP/1.1 Connection: Upgrade Upgrade: websocket Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ== ...

Сервер отвечает примерно вот так:

HTTP/1.1 101 Switching Protocols Upgrade: websocket Connection: Upgrade Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= ...

А дальше соединение не закрывается и между сервером и клиентом начинается двусторонний обмен сообщениями. Клиент и сервер в любой момент могут отправить по открытому соединению сообщение, а другая сторона его примет и обработает, но не ответит, а просто отправит сообщение сама тогда, когда это необходимо. В общем это очень полезная и эффективная технология, которая позволяет реализовать асинхронное взаимодействие с сервером.

Framing

После того как соединение установлено клиент и сервер отправляют друг другу так называемые «фреймы». Это сообщения со специально структурой. Подробно о них можно прочитать в спецификации. Но если коротко, то они состоят из:

[metadata][length][masking-key][payload-data]

• metadata — некая информация о сообщении и резервные биты

• length — длина сообщения

• masking-key — ключ для размаскирования

• payload-data — сами данные

То есть данные сообщения приходят не в чистом виде, а специальным образом измененные при помощи ключа. Причем клиент обязан всегда посылать маскированные данные, а сервер может отвечать как чистыми данными, так и маскированными, о чем он сообщает в мета-данных фрейма.

Нас интересует сам способ маскирования!

Masking/unmasking

Есть ключ — это 4 случайных байта. Нужно взять его и пройтись по всему массиву данных операцией побитового исключения. Допустим вот ключ:

key = {0, 1, 2, 3}; 

Вот данные:

payload = {0, 1, 2, 3, 4, 5}

Вот что делает функция маскирования (^ — побитовое исключение):

k ^ p == m 0 ^ 0 == 0 1 ^ 1 == 0 2 ^ 2 == 0 3 ^ 3 == 0 0 ^ 4 == 4 1 ^ 5 == 4

То есть вычисляет побитовое исключение каждого байта ключа с каждым байтом данных. Затем, дойдя до 5-го байта данных, алгоритм перемещает начало ключа на эту позицию и повторяет.

Проблема

А в чем проблема-то? Дело в том, что мы (мы — это группа из двух разработчиков WLJS Notebook) в своем проекте используем язык Wolfram. А он, как известно тем, кто этот язык не использует, «очень медленный», так как это символьный и интерпретируемый язык. И когда мы реализовывали протокол WebSocket для Wolfram Language (WL), то в полной мере это почувствовали. Дело в том, что нам требуется максимально быстрая скорость ответа сервера по протоколу WebSocket, но самая первая наивная реализация маскирования была медленная и тормозила взаимодействие браузера с сервером. Поэтому я занялся поиском путей решения этой проблемы. Этот путь я и изложу ниже.

Стандартная реализация

Итак, у нас язык Wolfram и алгоритм маскирования на нем выглядит вот так:

mask1 = Function[{maskingKey, payload},    ByteArray[Table[     BitXor[payload[[i]], maskingKey[[Mod[i - 1, 4] + 1]]],      {i, 1, Length[payload]}   ]] ]; 

• Function[{args}, expr] — чистая функция

• ByteArray[] — мы должны вернуть массив байт для записи в сокет

• Table[f[i], {i, 1, n}] — создает список с помощью цикла

• BitXor[a, b] — побитовое исключение

• Mod[a, b] — остаток от деления

Создадим один ключ:

maskingKey = ByteArray[{192, 168, 13, 7}]

И набор данных:

payload = StringToByteArray["{\"type\": \"point\", \"x\": 1.23, \"y\": 4.56}"]

Применим функцию маскирования:

payloadMasked = mask1[maskingKey, payload] (* => ByteArray[..] *)

В результате мы передали на вход два массива байт и получили на выходе один массив. Отобразить его в виде строки не получится, так как он содержит символы, которые не используются в UTF-8. Но мы знаем, что это нечитаемый текст. Еще у маскирования есть одно важное свойство! Повторное применение восстанавливает данные, т.е. выполняет размаскирование. Это значит что клиент присылает данные, сервер их размаскирует, затем создает ответ и маскирует той же функцией и отправляет клиенту. Клиент в свою очередь восстанавливает их той же операцией, зная ключ из заголовка фрейма. Получается, что выполняется следующее равенство:

mask1[maskingKey, mask1[maskingKey, payload]] == payload

Быстродействие

Замеры проводятся на ноутбуке ASUS TUF с Windows 11 и процессором AMD Ryzen 4800H 2.90 GHz. Но нас конечно в первую очередь будет интересовать относительный прирост производительности.

А теперь проверим насколько быстра и эффективна первая наивная реализация. Запустим код несколько раз на данных длиной 1 Мб:

payload1Mb = ByteArray[RandomInteger[{0, 255}, 2^20]]

И замер времени:

t1 = RepeatedTiming[mask1[maskingKey, payload1Mb]]

Запомним этот результат. От него мы и будем отталкиваться в дальнейших замерах. Уважаемым читателям, может показаться, что 1.7 секунд на 1 мегабайт информации это какое-то издевательство. Все так и есть. И этому эффекту есть объяснение, если очень глубоко покопаться в устройстве языка. Я постараюсь ответить на вопрос какого почему таааааааааааааааааааааааааааааааааааааааааааааааааааааааааааааааак долго.

Чтобы понять почему такая чудовищная в плохом смысле скорость вообще возможна, нам нужно провести небольшой ликбез по устройству языка Wolfram.

• WL — это язык с неизменяемыми структурами данных.
  Это значит, что практически любое выражение, переменная, массив и все что угодно созданное в языке является константами. Всякий раз когда мы пытаемся, например, изменить значение в списке чисел — происходит создание нового списка. Т.е. это жутко неэффективно, но зато надежно и безопасно если писать код в функциональном стиле.

• Он функциональный и символьный одновременно.
  Функциональный — это значит что он придерживается принципа, что есть данные и некие алгоритмы преобразования, которые возвращают данные. А символьный он по принципу исполнения кода на основе правил и замен, что тоже является очень неэффективной концепцией с точки зрения производительности, но супер крутой штукой концептуально. Подробнее про символьное программирование в WL можно почитать вот здесь.

• Он интерпретируемый.
  Это главная проблема. Ведь каждая атомарная операция проходит долгий путь прежде чем выполниться на процессоре. Вот как этот путь выглядит:

[user] -> [ui] -> [interpreter] -> [byte-code] -> [machine code] -> [cpu]

т.е. сначала пользователь вводит строку, на UI она выглядит как аналог HTML-разметки в виде так называемых «боксов», UI разбирает боксы и вытаскивает из них выражение, выражение интерпретатором превращается в полную форму, полная форма компилируется в байт-код виртуальной машины ядра, байт-код отправляется на ядро, ядро исполняет байт код превращая его в машинный код и отправляя на процессор. Затем эта цепочка выполняется в обратную сторону и заканчивается все тем, что фотоны/электромагнитные волны испущенные экраном компьютера попадают на сетчатку глаза и нервные импульсы двигаются в мозг передавая ответ ядра пользователю. То есть вся эта длинная цепочка выполняется на каждое атомарное действие. В нашем коде первый этап правда выполнится один раз, но затем каждый шаг цикла будет проходить через байт код и машинный код, а затем обратно. Сами по себе операции на CPU выполняются быстро, но если их очень много то эффективность сильно падает.

Еще одним важным фактором является то, что мы записали в наивной версии все очень просто и мы не обращали внимание на типы данных. На них тоже нужно остановиться. Как мы понимаем из условия на вход передаются массивы байт. Массив байт — это сложный композитный объект. На абсолютно каждой итерации цикла это объект целиком прогоняется до процессора и обратно, где у массива извлекается i-ый элемент. Что конечно же абсолютно неэффективно. Дело в том, что из-за сложной структуры массива байт, язык не способен применить одну из своих технологий, которая называется «автокомпиляция». Она срабатывает для численных алгоритмов, в тех случаях когда это простой цикл с численными операциями. Тогда ядро понимает, что ему не нужно каждый переход от итерации к итерации гонять через виртуальную машину и обратно. А выполняет итерации по циклу сразу в виртуальной машине, что сокращает оверхед. Но как этого добиться?

Автокомпиляция

Чтобы заставить ядро использовать эту технологию мы должны избавиться от того, что мешает этот код ускорить. У нас должны быть на входе два списка чисел и на выходе список чисел. Значит массив байт надо превратить в список перед циклом вот так:

mask2 = Function[{maskingKey, payload},  Module[{k = Normal[maskingKey], p = Normal[payload]},  ByteArray[Table[ BitXor[p[[i]], k[[Mod[i - 1, 4] + 1]]],  {i, 1, Length[p]} ]] ]];   t2 = RepeatedTiming[mask2[maskingKey, payload1Mb]]

В примере выше внутри цикла уже не осталось сложного объекта ByteArray, который мешал применить автокомпиляцию. Именно этим объясняется такое ускорение сразу в 23 раза. Важно помнить этой свойство языка и правильно работать с типами.

Compile

Если WL умеет самостоятельно применять автокомпиляцию — т.е. отбрасывать промежуточные шаги при обращении к ядру, то может быть можно и в ручную заставить язык это сделать? Конечно, да! Именно для этого предназначена функция Compile. Она создает байт-код для ядра, который не выполняется построчно, а весь исполняется в виртуальной машине. Вот как можно переделать код выше:

cmask3 = Compile[{{k, _Integer, 1}, {p, _Integer, 1}},  Table[ BitXor[p[[i]], k[[Mod[i - 1, 4] + 1]]],  {i, 1, Length[p]} ] ];   mask3[maskingKey_ByteArray, payload_ByteArray] :=  Module[{k = Normal[maskingKey], p = Normal[payload]},  ByteArray[cmask3[k, p]] ]  t3 = RepeatedTiming[mask3[maskingKey, payload1Mb]]

И снова внимательно посмотрим на типы данных. На вход передаются массивы байт. Специальная структура данных для хранения чисел с типом «UnsignedInteger8». Она хранит их наиболее эффективно. Но когда мы вызываем функцию Normal, то массив байт превращается в список целых чисел с типом long ли Integer64. То есть размер массива увеличивается в 8 раз, а эффективность атомарных операцией на 64-битных числах падает. Поему мы используем long? Увы, таковы ограничения Compile. Эта функция может работать только с 64-битными целыми, действительными машинной точности, комплексными числами и булевыми значениями True|False. Кроме того Compile работает с тензорами из перечисленных типов. Т.е. передать туда массив из чисел 8-битной точности невозможно. Есть ли способ это сделать другим путем? Конечно, да!

FunctionCompile

Начиная с 12 версии в языке появилась функция FunctionCompile. Это функция делает примерно тоже самое, что и Compile, т.е. превращает код на языке Wolfram в байт-код. Но если Compile возвращает функцию для вызывает байт-кода на виртуальной машине WL, то FunctionCompile компилирует IR-код для LLVM. Это самое важное отличие в концептуальном плане нового способа компиляции от старого. Еще одним отличием является то, что FunctionCompile поддерживает гораздо больше типов данных. И одним из них является как раз нужный нам массив байт. Что ж, возьмем тот же алгоритм и перепишем его вот так:

cmask4 = FunctionCompile[Function[ { Typed[maskingKey, "NumericArray"::["MachineInteger", 1]],  Typed[payload, "NumericArray"::["MachineInteger", 1]] },  Table[ BitXor[payload[[i]], maskingKey[[Mod[i - 1, 4] + 1]]],  {i, 1, Length[payload]}     ] ]];   mask4[maskingKey_ByteArray, payload1Mb_ByteArray] :=  ByteArray[cmask4[maskingKey, payload1Mb]]  t4 = RepeatedTiming[mask4[maskingKey, payload1Mb]]

Синтаксис очень похож на Compile, только теперь нужно использовать двойную конструкцию из FunctionCompile[Function[]], а типы аргументов следует указывать явно. Важно, что такая скомпилированная функция сразу может принимать на вход массивы байт без дополнительных преобразований. Но возвращает по прежнему список, который мы приводим к типу ByteArray[].

Опять же, внимательный читатель заметит, что я снова использовал 64-битные целые, вместо беззнаковых 8-битных. Все потому, что FunctionCompile в таком виде, как записано выше по умолчанию выводит тип, который возвращает Table как «PackedArray»::[«MachineInteger», 1]. Т.е. это упакованный обычный список. Можно ли заставить Table явно возвращать массив байт? Увы, ответ — нет. Внутри скомпилированного кода функция Table может вернуть только упакованный список. Но можно переписать код с использованием другого цикла, как это принято делать в некоторых других языках. Сначала мы создаем копию входного массива, а затем в цикле изменяем его значения выполняя побитовое исключение. Почему не менять сам входной массив? Потому что входные параметры неизменяемы и это сделать невозможно внутри FunctionCompile. Вот как выглядит новый код:

dec5 = FunctionDeclaration[ fmask5,  Typed[{ "NumericArray"::["UnsignedInteger8", 1],  "NumericArray"::["UnsignedInteger8", 1] } -> "NumericArray"::["UnsignedInteger8", 1]] @  Function[{maskingKey, payload},  Module[{result = payload},  Do[ result[[i]] =  BitXor[payload[[i]], maskingKey[[Mod[i - 1, 4] + 1]]],  {i, 1, Length[payload]} ];  result ] ] ];   cmask5 = FunctionCompile[dec5, fmask5];   mask5[maskingKey_ByteArray, payload1Mb_ByteArray] :=  ByteArray[cmask5[maskingKey, payload1Mb]]

• FunctionDeclaration[name, dec] — декларирует сигнатуру функции с точным указанием типов аргументов и возвращаемого результата. Затем задекларированную функцию можно скомпилировать.

• result = payload — эта строчка копирует исходный массив. Именно поэтому значения в result теперь можно менять.

Проверим быстродействие:

t5 = RepeatedTiming[mask5[maskingKey, payload1Mb]]

Что мы видим? Все эти накладные расходы на копирование, преобразования, передачу по значению и изменение значений в копии напрямую сделали код только медленнее! Но зато он эффективнее по использованию памяти.

У этого способа есть еще один минус. Не всякий пользователь сможет эту функцию скомпилировать из-за того что технология появилась относительно недавно и в версиях языка 12, 13 и 14 ведет себя по разному. Где-то код выше в принципе не скомпилируется. Плюс сама по себе компиляция занимает много времени. Конечно результат можно сохранить на диск, но это нужно сделать для всех поддерживаемых операционных систем. В итоге накладные расходы и сложность поддержки экспериментальной функции имею свои минусы и я продолжил поиск путей оптимизации.

LibraryLink

Как известно, ядро Wolfram Language написано на Си. Плюс многие критически важные функции тоже реализованы на Си. Можно ли самостоятельно загрузить в ядро такую функцию? Конечно да, это делается при помощи встроенной библиотеки LibraryLink. Она предоставляет функции для компиляции и использования кода на языке Си прямо в WL.

Как выглядит работа с ней:

• Пишем код на Си

• Оборачиваем его в специальные инструкции из библиотеки «WolframLibrary.h»

• На стороне WL компилируем при помощи CreateLibrary

• И загружаем при помощи LibraryFunctionLoad

• Готово!

Но лучше один раз увидеть… Далее будет часть кода на Си.

В каждой библиотеке на языке Си необходимо всегда выполнять ритуал импорта вот такой:

#include "WolframLibrary.h" #include "WolframNumericArrayLibrary.h"

Затем для ядра WL необходимо определить три функции, которые инициализируют библиотеку, деинициализируют и возвращают версию LibraryLink, под которую библиотека была скомпилирована:

DLLEXPORT mint WolframLibrary_getVersion() {     return WolframLibraryVersion; }  DLLEXPORT int WolframLibrary_initialize(WolframLibraryData libData) {     return LIBRARY_NO_ERROR; }  DLLEXPORT void WolframLibrary_uninitialize(WolframLibraryData libData) {     return; }

Есть еще небольшая тонкость для работы этой библиотеки на всех ОС. Надо добавить зависимости и одну инструкцию вот так:

#include <stdio.h> #include <stdlib.h> #include <string.h>  #define BYTE uint8_t

Все готово. Теперь приступим к написанию основного кода.

Первое, что я сделал — это переписал алгоритм маскирования в лоб. Вот как выглядит функция, которая его реализует:

DLLEXPORT int mask6(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {     //get payload and length     MNumericArray nMask = MArgument_getMNumericArray(Args[0]);     BYTE* mask = (BYTE *)libData->numericarrayLibraryFunctions->MNumericArray_getData(nMask);     mint maskLen = MArgument_getInteger(Args[1]);          //get maskingKey and length     MNumericArray nArr = MArgument_getMNumericArray(Args[2]);     BYTE* arr = (BYTE *)libData->numericarrayLibraryFunctions->MNumericArray_getData(nArr);      mint arrLen = MArgument_getInteger(Args[3]);      //create new result array     BYTE *result = (BYTE*)malloc(sizeof(BYTE) * arrLen);       //bit xor     for (size_t i = 0; i < arrLen; i++) {         result[i] = arr[i] ^ mask[i % maskLen];      }          //creating a struct for sending to WL side     MNumericArray nResult;     libData->numericarrayLibraryFunctions->MNumericArray_new(MNumericArray_Type_UBit8, 1, &arrLen, &nResult);     memcpy((uint8_t*)libData->numericarrayLibraryFunctions->MNumericArray_getData(nResult), result, arrLen);      //return result and free memory     MArgument_setMNumericArray(Res, nResult);     free(result);      return LIBRARY_NO_ERROR; }

Важно отметить, что:

• Все функции подключаемые к ядру WL должны иметь аргумент и возвращаемые результат как указано выше.

• libData — структура с набором функций для связи с ядром

• Argc — число аргументов на стороне WL

• Args — сами аргументы

• Res — структура в которую записывается результат — то, что вернется на стороне WL

• На стороне Си возвращается код ошибки

Вернемся обратно в блокнот и скомпилируем этот исходник вот так:

Get["CCompilerDriver`"];  Get["LibraryLink`"];   lib = CreateLibrary[File["internal.c"], "internal"]; 

В lib сохранился путь к dll-файлу с машинным кодом, специфичным для ОС. Загрузить из него созданную функцию можно вот так:

cmask6 = LibraryFunctionLoad[ lib,  "mask6",  {"ByteArray", Integer, "ByteArray", Integer},  "ByteArray" ]

А затем создаем более удобное определение, вид которого совпадет с предыдущими вариантами:

mask6[maskingKey_ByteArray, payload_ByteArray] :=  cmask6[ maskingKey, Length[maskingKey],  payload, Length[payload] ]

И измерим время:

t6 = RepeatedTiming[mask6[maskingKey, payload1Mb]]

Да, действительно эффективность снизилась значительно сильно по сравнению с двумя последними вариантами. Но зато теперь мы можем без особых сложностей и преобразований передавать массив байт и возвращать массив байт. Далее необходимо очень внимательно посмотреть на код и понять как теперь улучшить сам алгоритм!

Я приведу сразу весь список улучшений, которые можно сделать:

• Можно избавится от i % maskLen используя двойной цикл

int k;  int len = arrLen - arrLen % maskLen;  for (size_t i = 0; i < len; i = i + maskLen) {   for (size_t j = 0; j < maskLen; j++) {     k = i + j;      result[k] = arr[k] ^ mask[j];    } }  for (size_t i = len; i < arrLen; i++) {   result[i] = arr[i] ^ mask[i % maskLen]; }  /* вместо for (size_t i = 0; i < arrLen; i++) {   result[i] = arr[i] ^ mask[i % maskLen];  }*/

Кажется, что операций стало больше, но хвост почти не занимает времени, а вот замена одного оператора остатка от деления на одно присвоение и одно сложение дает очень существенную разницу.

• Далее если внимательно изучить руководство по LibraryLink, то можно понять, что создание result — лишняя операция и тем более копирование результата в итоговый массив. Дело в том, что у любого созданного массива (MNumericArray()) можно извлечь указатель на сам участок памяти где хранятся сырые данные в виде массива чисел. Поэтому можно его и использовать в качестве result:

MNumericArray nResult; libData->numericarrayLibraryFunctions->MNumericArray_new(MNumericArray_Type_UBit8, 1, &arrLen, &nResult); BYTE *result = (BYTE*)libData->numericarrayLibraryFunctions->MNumericArray_getData(nResult);   /* вместо BYTE *result = (BYTE*)malloc(sizeof(BYTE) * arrLen);  MNumericArray nResult; libData->numericarrayLibraryFunctions->MNumericArray_new(MNumericArray_Type_UBit8, 1, &arrLen, &nResult); memcpy((uint8_t*)libData->numericarrayLibraryFunctions->MNumericArray_getData(nResult), result, arrLen); free(result);*/

Итоговый код на Си будет вот такой:

DLLEXPORT int mask7(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {     MNumericArray nMask = MArgument_getMNumericArray(Args[0]);     BYTE* mask = (BYTE *)libData->numericarrayLibraryFunctions->MNumericArray_getData(nMask);     mint maskLen = MArgument_getInteger(Args[1]);          MNumericArray nArr = MArgument_getMNumericArray(Args[2]);     BYTE* arr = (BYTE *)libData->numericarrayLibraryFunctions->MNumericArray_getData(nArr);      mint arrLen = MArgument_getInteger(Args[3]);      MNumericArray nResult;     libData->numericarrayLibraryFunctions->MNumericArray_new(MNumericArray_Type_UBit8, 1, &arrLen, &nResult);     BYTE *result = (BYTE*)libData->numericarrayLibraryFunctions->MNumericArray_getData(nResult);       const int len = arrLen - arrLen % maskLen;      int k;       for (size_t i = 0; i < len; i = i + maskLen) {         for (size_t j = 0; j < maskLen; j++) {             k = i + j;              result[k] = arr[k] ^ mask[j];          }     }      for (size_t i = len; i < arrLen; i++) {         result[i] = arr[i] ^ mask[i % maskLen];     }      MArgument_setMNumericArray(Res, nResult);     return LIBRARY_NO_ERROR; }

• И последнее, но уже на стороне WL. Каждый раз при вызове библиотечной функции массив байт будет копировать из памяти, которая доступна ядру в ту память, которая доступна библиотечной функции. С одной стороны это более безопасно, так как даже если внутри библиотечной функции изменить входной массив, то это не повлияет на то, что происходит в сессии ядра. С другой стороны — копирование занимает время. Но у LibraryLink есть механизм использования общей памяти для аргументов. Это значит, что копирования не будет, а будет передача ссылки на уже существующий в сессии массив. Платой за это является то, что попытки поменять значения в общем массиве или очистить его приводят к падению ядра. В общем использовать этот механизм можно вот так:

cmask7 = LibraryFunctionLoad[   lib,    "mask7",    {     "ByteArray",      Integer,      {LibraryDataType["ByteArray"], "Shared"},      Integer   },    "ByteArray" ]

Ну а далее все как обычно:

mask7[maskingKey_ByteArray, payload_ByteArray] :=  cmask7[ maskingKey, Length[maskingKey],  payload, Length[payload] ]  t7 = RepeatedTiming[mask7[maskingKey, payload1Mb]]

Посмотрим, что получилось!

Итоговая разница в скорости выполнения с самой первой наивной реализацией в 2550 раз, а с самой быстрой предыдущей реализацией, скомпилированной для LLVM в 5.8 раз. В итоге теперь данные размером в 1 Мб маскируются и демаскируются примерно за 666 мкс, что как бы намекает…

Выводы

Еще раз повторюсь, что в первую очередь пытался продемонстрировать подход к написанию эффективного кода для Wolfram Language. Да, в итоге я обратился к Си, но и этот способ имеет право на жизнь и часто нами используется в различных библиотеках. Подтвердилось ли, что WL — медленный язык? И да и нет. Если писать код бездумно — то уже ничего не поможет, если же погрузиться в WL дальше уровня решения уравнений и взятия интегралов и использовать его как язык программирования — то всегда можно найти эффективное решение для многих задач. Всем спасибо за внимание!