Создаём быстрые gRPC-сервисы с Tonic и Rust

a7f9c7bd490b0c32ee4b4b65413599f3.png

Привет, Хабр!

Сегодня посмотрим, как с помощью фреймворка Tonic и языка Rust создавать gRPC-сервисы для задач машинного обучения. Если в вашем проекте нужно максимально эффективно строить распределённые системы, а производительность и асинхронное программирование — это то, что вы цените, то Rust в связке с Tonic станет отличным инструментом

Установка Tonic

Первое, что нужно сделать, это создать новый проект на Rust. Для этого открываем терминал и вводим:

cargo new my_grpc_service
cd my_grpc_service

Теперь открываем файл Cargo.toml. Добавим Tonic и некоторые другие необходимые библиотеки.

[dependencies]
tonic = "0.12.2"
prost = "0.11"
tokio = { version = "1", features = ["full"] }
  • Tonic — основной фреймворк для gRPC.

  • Prost — библиотека для кодирования и декодирования Protocol Buffers.

  • Tokio — асинхронный runtime для Rust, который нам нужен для работы с Tonic.

Теперь добавим необходимые флаги функций для Tonic. Это нужно для включения всех необходимых возможностей. В Cargo.toml должно быть что-то вроде этого:

[dependencies]
tonic = { version = "0.12.2", features = ["transport", "codegen", "tls"] }
prost = "0.11"
tokio = { version = "1", features = ["full"] }

[build-dependencies]
tonic-build = "0.12.2"
  • transport — включает поддержку HTTP/2.

  • codegen — включает кодогенерацию из .proto файлов.

  • tls — для работы с безопасными соединениями.

Теперь, когда есть зависимости, создадим базовую структуру нашего проекта.

mkdir proto

Создадим файл hello.proto в этой папке и добавим следующий код:

syntax = "proto3";

package hello;

service Greeter {
    rpc SayHello (HelloRequest) returns (HelloResponse);
}

message HelloRequest {
    string name = 1;
}

message HelloResponse {
    string message = 1;
}

Этот простой протокол описывает сервис Greeter, который принимает HelloRequest и возвращает HelloResponse.

Теперь нужно сгенерировать код из .proto файла. Для этого создаем файл build.rs в корне проекта с содержимым:

fn main() {
    tonic_build::compile_protos("proto/hello.proto").unwrap();
}

Этот код заставит tonic-build скомпилировать наш протокол при сборке проекта.

Теперь, когда у нас есть всё готово, пора собрать проект. Выполняем команду:

cargo build

Если всё сделано правильно, то проект скомпилируется без ошибок.

Создание gRPC-сервиса

Начнем с того, что нужно создать .proto файл, в котором определим наши сообщения и сервисы. Возьмем тот же hello.proto, который мы создали ранее, и немного расширим его. Допустим, нужно добавить функционал для приветствия пользователей и получения списка пользователей:

syntax = "proto3";

package hello;

service Greeter {
    rpc SayHello (HelloRequest) returns (HelloResponse);
    rpc ListUsers (ListUsersRequest) returns (ListUsersResponse);
}

message HelloRequest {
    string name = 1;
}

message HelloResponse {
    string message = 1;
}

message ListUsersRequest {}

message ListUsersResponse {
    repeated string users = 1;
}

Здесь добавили новый метод ListUsers, который возвращает список пользователей. Теперь сгенерируем код.

Выполняем команду сборки:

cargo build

Это сгенерирует Rust-модули на основе .proto файла.

Теперь создаем файл src/main.rs и добавляем следующий код:

use tonic::{transport::Server, Request, Response, Status};
use hello::greeter_server::{Greeter, GreeterServer};
use hello::{HelloRequest, HelloResponse, ListUsersRequest, ListUsersResponse};

pub mod hello {
    include!("generated.hello.rs");
}

#[derive(Default)]
pub struct MyGreeter {}

#[tonic::async_trait]
impl Greeter for MyGreeter {
    async fn say_hello(
        &self,
        request: Request,
    ) -> Result, Status> {
        let name = request.into_inner().name;
        let reply = HelloResponse {
            message: format!("Hello, {}!", name),
        };
        Ok(Response::new(reply))
    }

    async fn list_users(
        &self,
        _request: Request,
    ) -> Result, Status> {
        let users = vec!["Alice".to_string(), "Bob".to_string(), "Charlie".to_string()];
        let reply = ListUsersResponse { users };
        Ok(Response::new(reply))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box> {
    let addr = "[::1]:50051".parse()?;
    let greeter = MyGreeter::default();

    Server::builder()
        .add_service(GreeterServer::new(greeter))
        .serve(addr)
        .await?;

    Ok(())
}

MyGreeter: реализация сервиса Greeter. Реализуем методы, определенные в нашем .proto файле.

say_hello: метод обрабатывает запрос на приветствие. Он извлекает имя из HelloRequest, формирует ответ и возвращает его.

list_users: метод возвращает список пользователей. Просто создаем вектор со строками и оборачиваем его в ListUsersResponse.

В Tonic обработка запросов и ответов происходит с помощью типов Request и Response. Эти типы позволяют нам работать с данными, поступающими от клиента, и формировать ответы, которые будут отправляться обратно.

Request: Этот тип содержит данные, полученные от клиента, а также метаданные. Можно извлечь основное сообщение с помощью метода into_inner().

let name = request.into_inner().name;

Response: Этот тип используется для формирования ответов. Создаем экземпляр Response, передавая в него данные, которые хотите вернуть клиенту.

let reply = HelloResponse {
    message: format!("Hello, {}!", name),
};
Ok(Response::new(reply))

Чтобы запустить сервер, просто выполняем команду:

cargo run

Если всё прошло успешно, вы увидите, что сервер запущен на localhost:50051.

Немного оптимизации

Управление размером сообщений

Первый шаг к оптимизации — это управление размерами сообщений.

В Tonic есть возможность задавать максимальные размеры входящих и исходящих сообщений. По дефолту лимит составляет 4 МБ для входящих сообщений и usize::MAX для исходящих.

Добавим следующую конфигурацию в реализацию сервера:

use tonic::{transport::Server, Request, Response, Status};
use hello::greeter_server::{Greeter, GreeterServer};
use hello::{HelloRequest, HelloResponse};

#[tokio::main]
async fn main() -> Result<(), Box> {
    let addr = "[::1]:50051".parse()?;
    let greeter = MyGreeter::default();

    Server::builder()
        .max_decoding_message_size(10 * 1024 * 1024) // 10 МБ
        .max_encoding_message_size(10 * 1024 * 1024) // 10 МБ
        .add_service(GreeterServer::new(greeter))
        .serve(addr)
        .await?;

    Ok(())
}

Здесь установили пределы в 10 МБ для входящих и исходящих сообщений.

Сжатие — ещё один мощный инструмент оптимизации. Tonic поддерживает сжатие сообщений с использованием алгоритмов gzip и zstd.

Для начала нужно будет включить сжатие в конфигурации сервера. Вот как это можно сделать:

use tonic::{transport::Server, Request, Response, Status};
use hello::greeter_server::{Greeter, GreeterServer};
use hello::{HelloRequest, HelloResponse};

#[tokio::main]
async fn main() -> Result<(), Box> {
    let addr = "[::1]:50051".parse()?;
    let greeter = MyGreeter::default();

    let compression = tonic::Compression::Gzip;

    Server::builder()
        .compression(compression)
        .add_service(GreeterServer::new(greeter))
        .serve(addr)
        .await?;

    Ok(())
}

Теперь сервер будет использовать сжатие gzip для всех сообщений. Если нужно использовать zstd, просто заменяем tonic::Compression::Gzip на tonic::Compression::Zstd.

На стороне клиента можно включить сжатие:

let mut client = GreeterClient::connect("http://[::1]:50051")
    .await?
    .accept_compressed(tonic::Compression::Gzip)
    .await?;

Таким образом, при отправке данных они будут автоматически сжиматься, а при получении — распаковываться.

Tonic построен на основе tokio, что позволяет использовать асинхронные потоки для обработки запросов.

Допустим, нужно обрабатывать запросы с задержкой, чтобы эмулировать взаимодействие с удаленной БД или API. Можно использовать tokio::time::sleep для создания задержки:

#[tonic::async_trait]
impl Greeter for MyGreeter {
    async fn say_hello(
        &self,
        request: Request,
    ) -> Result, Status> {
        let name = request.into_inner().name;

        // Имитация задержки
        tokio::time::sleep(std::time::Duration::from_secs(2)).await;

        let reply = HelloResponse {
            message: format!("Hello, {}!", name),
        };
        Ok(Response::new(reply))
    }
}

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

Подробнее с Tonic можно ознакомиться здесь.

Изучить продвинутые ML приемы для практикующих Data Scientists,  желающих повысить свой профессиональный уровень до Middle+, можно на онлайн-курсе «Machine Learning. Advanced».

© Habrahabr.ru