| 1465 | } |
| 1466 | |
| 1467 | void ProxyRunner::forward_query(TcpClient::ConnectionId client_connection_id, |
| 1468 | ton::tl_object_ptr<cocoon_api::client_runQueryEx> req) { |
| 1469 | auto client_tcp_connection = static_cast<ProxyInboundClientConnection *>(get_connection(client_connection_id)); |
| 1470 | auto client = static_cast<ProxyInboundClientConnection *>(client_tcp_connection)->client_info(); |
| 1471 | |
| 1472 | auto client_request_id = req->request_id_; |
| 1473 | auto fail = [&](td::Status S) { |
| 1474 | td::BufferSlice res; |
| 1475 | auto client_proto_version = client_tcp_connection->proto_version(); |
| 1476 | if (client_proto_version > 0) { |
| 1477 | res = cocoon::create_serialize_tl_object<cocoon_api::client_queryAnswerErrorEx>( |
| 1478 | client_request_id, S.code(), S.message().str(), 1, |
| 1479 | ton::create_tl_object<cocoon_api::client_queryFinalInfo>( |
| 1480 | (client_proto_version >= 2 ? 2 : 0), ton::create_tl_object<cocoon_api::tokensUsed>(0, 0, 0, 0, 0), "", "", |
| 1481 | td::Clocks::system(), td::Clocks::system(), td::Clocks::system(), td::Clocks::system())); |
| 1482 | } else { |
| 1483 | res = cocoon::create_serialize_tl_object<cocoon_api::client_queryAnswerError>( |
| 1484 | S.code(), S.message().str(), client_request_id, ton::create_tl_object<cocoon_api::tokensUsed>(0, 0, 0, 0, 0)); |
| 1485 | } |
| 1486 | send_message_to_connection(client_connection_id, std::move(res)); |
| 1487 | stats_->requests_rejected++; |
| 1488 | }; |
| 1489 | |
| 1490 | if (client->tokens_available() < req->max_tokens_) { |
| 1491 | return fail(td::Status::Error( |
| 1492 | ton::ErrorCode::error, |
| 1493 | PSTRING() << "client balance is too low: max_tokens=" << req->max_tokens_ |
| 1494 | << " tokens_payed=" << client->tokens_payed() << " tokens_used=" << client->tokens_used() |
| 1495 | << " tokens_reserved=" << client->tokens_reserved() << " tokens_stake=" << client->tokens_stake())); |
| 1496 | } |
| 1497 | |
| 1498 | if ((td::uint32)req->min_config_version_ > active_config_version_) { |
| 1499 | return fail(td::Status::Error(ton::ErrorCode::error, "active config version is too low")); |
| 1500 | } |
| 1501 | if (!client->allow_queries()) { |
| 1502 | return fail(td::Status::Error(ton::ErrorCode::notready, "client is closing")); |
| 1503 | } |
| 1504 | |
| 1505 | td::Bits256 encrypted_with = td::Bits256::zero(); |
| 1506 | if ((req->flags_ & 2) && !req->public_key_.is_zero()) { |
| 1507 | if (!get_private_key(req->public_key_, encrypted_with)) { |
| 1508 | return fail( |
| 1509 | td::Status::Error(ton::ErrorCode::error, PSTRING() << "unknown public key " << req->public_key_.to_hex())); |
| 1510 | } |
| 1511 | } |
| 1512 | |
| 1513 | auto to_reserve = adjust_tokens(req->max_tokens_ + req->query_.size(), req->max_coefficient_, 10000); |
| 1514 | auto R_worker_connection_id = |
| 1515 | choose_connection(req->model_name_, client->tokens_available(), (td::uint32)req->max_coefficient_, to_reserve, |
| 1516 | encrypted_with.is_zero() ? 0 : 3); |
| 1517 | if (R_worker_connection_id.is_error()) { |
| 1518 | return fail(R_worker_connection_id.move_as_error()); |
| 1519 | } |
| 1520 | auto worker_connection = R_worker_connection_id.move_as_ok(); |
| 1521 | |
| 1522 | if (!client->reserve(to_reserve)) { |
| 1523 | return fail(td::Status::Error( |
| 1524 | ton::ErrorCode::error, |
nothing calls this directly
no test coverage detected