├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md └── squangle ├── base ├── Base.cpp ├── Base.h ├── ConnectionKey.cpp ├── ConnectionKey.h └── ExceptionUtil.h ├── logger ├── DBEventCounter.cpp ├── DBEventCounter.h ├── DBEventLogger.cpp └── DBEventLogger.h ├── mysql_client ├── AsyncConnectionPool.cpp ├── AsyncConnectionPool.h ├── AsyncHelpers.cpp ├── AsyncHelpers.h ├── AsyncMysqlClient.cpp ├── AsyncMysqlClient.h ├── ChangeUserOperation.cpp ├── ChangeUserOperation.h ├── ClientPool.h ├── Compression.cpp ├── Compression.h ├── ConnectOperation.cpp ├── ConnectOperation.h ├── ConnectPoolOperation.h ├── Connection.cpp ├── Connection.h ├── ConnectionHolder.cpp ├── ConnectionHolder.h ├── ConnectionOptions.cpp ├── ConnectionOptions.h ├── ConnectionPool.cpp ├── ConnectionPool.h ├── DbResult.cpp ├── DbResult.h ├── FetchOperation.cpp ├── FetchOperation.h ├── Flags.cpp ├── Flags.h ├── InternalConnection.cpp ├── InternalConnection.h ├── MultiQueryOperation.cpp ├── MultiQueryOperation.h ├── MultiQueryStreamOperation.cpp ├── MultiQueryStreamOperation.h ├── MysqlClientBase.cpp ├── MysqlClientBase.h ├── MysqlExceptionBuilder.h ├── Operation.cpp ├── Operation.h ├── OperationBatch.cpp ├── OperationBatch.h ├── OperationHelpers.cpp ├── OperationHelpers.h ├── PoolKey.h ├── PoolStorage.h ├── Query.cpp ├── Query.h ├── QueryBenchmark.cpp ├── QueryGenerator.h ├── QueryOperation.cpp ├── QueryOperation.h ├── README.md ├── ResetOperation.cpp ├── ResetOperation.h ├── Row.cpp ├── Row.h ├── SSLOptionsProviderBase.cpp ├── SSLOptionsProviderBase.h ├── SemiFutureAdapter.cpp ├── SemiFutureAdapter.h ├── SpecialOperation.cpp ├── SpecialOperation.h ├── SyncConnectionPool.cpp ├── SyncConnectionPool.h ├── SyncMysqlClient.cpp ├── SyncMysqlClient.h ├── TODO ├── TwoLevelCache.h └── mysql_protocol │ ├── MysqlConnectOperationImpl.cpp │ ├── MysqlConnectOperationImpl.h │ ├── MysqlConnectPoolOperationImpl.h │ ├── MysqlConnection.cpp │ ├── MysqlConnection.h │ ├── MysqlFetchOperationImpl.cpp │ ├── MysqlFetchOperationImpl.h │ ├── MysqlOperationImpl.cpp │ ├── MysqlOperationImpl.h │ ├── MysqlResult.cpp │ ├── MysqlResult.h │ ├── MysqlRow.h │ ├── MysqlRowMetadata.h │ ├── MysqlSpecialOperationImpl.cpp │ └── MysqlSpecialOperationImpl.h └── util ├── StorageRow.cpp ├── StorageRow.h ├── StringStore.h └── tests └── StorageRowTest.cpp /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to make participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at . All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 72 | 73 | [homepage]: https://www.contributor-covenant.org 74 | 75 | For answers to common questions about this code of conduct, see 76 | https://www.contributor-covenant.org/faq 77 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to SQuangLe 2 | 3 | ## Pull Requests 4 | We actively welcome your pull requests. 5 | 1. Fork the repo and create your branch from `main`. 6 | 2. If you've added code that should be tested, add tests 7 | 3. If you've changed APIs, update the documentation. 8 | 4. Ensure the test suite passes. 9 | 5. Make sure your code lints. 10 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 11 | 12 | ## Contributor License Agreement ("CLA") 13 | In order to accept your pull request, we need you to submit a CLA. You only need 14 | to do this once to work on any of Facebook's open source projects. 15 | 16 | Complete your CLA here: 17 | 18 | Pull requests will not be directly merged into this repository - they will be 19 | mirrored into Facebook's internal repository, then pushed to github, along 20 | with any other changes. 21 | 22 | ## Issues 23 | 24 | Squangle is currently only supported as a dependency of HHVM, not as a 25 | standalone project. 26 | 27 | Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe 28 | disclosure of security bugs. In those cases, please go through the process 29 | outlined on that page and do not file a public issue. 30 | 31 | ## License 32 | By contributing to SQuangLe, you agree that your contributions will be 33 | licensed under its BSD license. 34 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD License 2 | 3 | For Squangle software 4 | 5 | Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. 6 | 7 | Redistribution and use in source and binary forms, with or without modification, 8 | are permitted provided that the following conditions are met: 9 | 10 | * Redistributions of source code must retain the above copyright notice, this 11 | list of conditions and the following disclaimer. 12 | 13 | * Redistributions in binary form must reproduce the above copyright notice, 14 | this list of conditions and the following disclaimer in the documentation 15 | and/or other materials provided with the distribution. 16 | 17 | * Neither the name Facebook nor the names of its contributors may be used to 18 | endorse or promote products derived from this software without specific 19 | prior written permission. 20 | 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 22 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 23 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 25 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 26 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 28 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 30 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | SQuangLe 2 | ======== 3 | 4 | Overview 5 | -------- 6 | 7 | SQuangLe is a C++ MySQL client library built on top of 8 | [WebScaleSQL](http://webscalesql.org/)'s C client library. It does not require 9 | a WebScaleSQL server - WebScaleSQL, MySQL, MariaDB, and Percona Server should 10 | all work fine. 11 | 12 | Current Status 13 | -------------- 14 | 15 | SQuangLe is not supported as a standalone project; it is released as a 16 | dependency of HHVM's async mysql extension to the Hack language. 17 | 18 | Features 19 | -------- 20 | 21 | - Object-oriented API 22 | - Query builder with automatic escaping 23 | - Asynchronous query execution 24 | 25 | License 26 | ------- 27 | 28 | SQuangLe is BSD-licensed. We also provide an additional patent grant. Please 29 | see the LICENSE file. 30 | 31 | Contributing 32 | ------------ 33 | 34 | Please see CONTRIBUTING.md 35 | -------------------------------------------------------------------------------- /squangle/base/Base.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | 11 | #include "squangle/base/Base.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | std::string_view toString(QueryCallbackReason reason) { 16 | switch (reason) { 17 | case QueryCallbackReason::RowsFetched: 18 | return "RowsFetched"; 19 | case QueryCallbackReason::QueryBoundary: 20 | return "QueryBoundary"; 21 | case QueryCallbackReason::Failure: 22 | return "Failure"; 23 | case QueryCallbackReason::Success: 24 | return "Success"; 25 | } 26 | 27 | LOG(DFATAL) << "unable to convert reason to string: " 28 | << static_cast(reason); 29 | return "Unknown reason"; 30 | } 31 | 32 | // overload of operator<< for QueryCallbackReason 33 | std::ostream& operator<<(std::ostream& os, QueryCallbackReason reason) { 34 | return os << toString(reason); 35 | } 36 | 37 | } // namespace facebook::common::mysql_client 38 | -------------------------------------------------------------------------------- /squangle/base/Base.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | 15 | namespace facebook::common::mysql_client { 16 | 17 | using Millis = std::chrono::milliseconds; 18 | using Micros = std::chrono::microseconds; 19 | using Duration = std::chrono::duration; 20 | using Timepoint = std::chrono::time_point; 21 | 22 | // The extra StringHasher and std::equal_to<> makes this unordered map 23 | // "transparent", meaning we can .find() by std::string, std::string_view, and 24 | // by const char*. */ 25 | using AttributeMap = folly::F14NodeMap; 26 | 27 | // For control flows in callbacks. This indicates the reason a callback was 28 | // fired. When a pack of rows if fetched it is used RowsFetched to 29 | // indicate that new rows are available. QueryBoundary means that the 30 | // fetching for current query has completed successfully, and if any 31 | // query failed (OperationResult is Failed) we use Failure. Success is for 32 | // indicating that all queries have been successfully fetched. 33 | enum class QueryCallbackReason { RowsFetched, QueryBoundary, Failure, Success }; 34 | 35 | // overload of operator<< for QueryCallbackReason 36 | std::ostream& operator<<(std::ostream& os, QueryCallbackReason reason); 37 | 38 | } // namespace facebook::common::mysql_client 39 | -------------------------------------------------------------------------------- /squangle/base/ConnectionKey.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/base/ConnectionKey.h" 10 | 11 | #include 12 | 13 | namespace facebook { 14 | namespace common { 15 | namespace mysql_client { 16 | 17 | MysqlConnectionKey::MysqlConnectionKey( 18 | folly::StringPiece host, 19 | int port, 20 | folly::StringPiece db_name, 21 | folly::StringPiece user, 22 | folly::StringPiece password, 23 | folly::StringPiece special_tag, 24 | bool ignore_db_name, 25 | folly::StringPiece unixSocketPath) 26 | : host_(host.toString()), 27 | dbName_(db_name.toString()), 28 | user_(user.toString()), 29 | password_(password.toString()), 30 | specialTag_(special_tag.toString()), 31 | unixSocketPath_(unixSocketPath), 32 | partialHash_(folly::Hash()(host, port, user, password, special_tag)), 33 | hash_( 34 | ignore_db_name ? partialHash_ : folly::Hash()(partialHash_, dbName_)), 35 | port_(port), 36 | ignoreDbName_(ignore_db_name) {} 37 | 38 | bool MysqlConnectionKey::operator==( 39 | const MysqlConnectionKey& rhs) const noexcept { 40 | return hash_ == rhs.hash_ && host_ == rhs.host_ && port_ == rhs.port_ && 41 | (ignoreDbName_ || dbName_ == rhs.dbName_) && user_ == rhs.user_ && 42 | password_ == rhs.password_ && specialTag_ == rhs.specialTag_ && 43 | unixSocketPath_ == rhs.unixSocketPath_; 44 | } 45 | 46 | bool MysqlConnectionKey::operator==(const ConnectionKey& rhs) const noexcept { 47 | try { 48 | const auto& key = dynamic_cast(rhs); 49 | return *this == key; 50 | } catch (const std::bad_cast& /*e*/) { 51 | // `rhs` was not a MysqlConnectionKey object 52 | return false; 53 | } 54 | } 55 | 56 | bool MysqlConnectionKey::partialEqual(const ConnectionKey& rhs) const noexcept { 57 | try { 58 | const auto& key = dynamic_cast(rhs); 59 | return partialHash_ == key.partialHash_ && host_ == key.host_ && 60 | port_ == key.port_ && user_ == key.user_ && 61 | password_ == key.password_ && specialTag_ == key.specialTag_ && 62 | unixSocketPath_ == key.unixSocketPath_; 63 | } catch (const std::bad_cast& /*e*/) { 64 | // `rhs` was not a MysqlConnectionKey object 65 | return false; 66 | } 67 | } 68 | 69 | std::string MysqlConnectionKey::getDisplayString(bool level2) const { 70 | if (unixSocketPath_.empty()) { 71 | return fmt::format( 72 | "{} [{}] ({}@{}:{})", 73 | level2 ? "" : dbName_, 74 | specialTag_, 75 | user_, 76 | host_, 77 | port_); 78 | } 79 | 80 | return fmt::format( 81 | "{} [{}] ({}@{})", 82 | level2 ? "" : dbName_, 83 | specialTag_, 84 | user_, 85 | unixSocketPath_); 86 | } 87 | } // namespace mysql_client 88 | } // namespace common 89 | } // namespace facebook 90 | -------------------------------------------------------------------------------- /squangle/base/ConnectionKey.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | // This class encapsulates the data that differentiates 2 connections: 16 | // host, port, db name and user. We also store the password to avoid 17 | // allowing a connection with wrong password be accepted. 18 | // We also store the key as string (without the password and special tag) 19 | // for debugging purposes and to use as keys in other maps 20 | class ConnectionKey { 21 | public: 22 | virtual ~ConnectionKey() = default; 23 | 24 | [[nodiscard]] virtual size_t hash() const = 0; 25 | 26 | [[nodiscard]] virtual size_t partial_hash() const = 0; 27 | 28 | [[nodiscard]] virtual bool partialEqual( 29 | const ConnectionKey& rhs) const noexcept = 0; 30 | 31 | [[nodiscard]] virtual std::string getDisplayString( 32 | bool level2 = false) const = 0; 33 | 34 | [[nodiscard]] virtual const std::string& host() const = 0; 35 | 36 | [[nodiscard]] virtual const std::string& db_name() const noexcept = 0; 37 | 38 | [[nodiscard]] virtual const std::string& user() const noexcept = 0; 39 | 40 | [[nodiscard]] virtual const std::string& password() const noexcept = 0; 41 | 42 | [[nodiscard]] virtual int port() const = 0; 43 | 44 | [[nodiscard]] virtual const std::string& special_tag() const noexcept = 0; 45 | 46 | [[nodiscard]] virtual bool operator==( 47 | const ConnectionKey& rhs) const noexcept = 0; 48 | 49 | [[nodiscard]] bool operator!=(const ConnectionKey& rhs) const noexcept { 50 | return !(*this == rhs); 51 | } 52 | }; 53 | 54 | class MysqlConnectionKey : public ConnectionKey { 55 | public: 56 | explicit MysqlConnectionKey( 57 | folly::StringPiece sp_host = "", 58 | int sp_port = 0, 59 | folly::StringPiece sp_db_name = "", 60 | folly::StringPiece sp_user = "", 61 | folly::StringPiece sp_password = "", 62 | folly::StringPiece sp_special_tag = "", 63 | bool sp_ignore_db_name = false, 64 | folly::StringPiece sp_unixSocketPath = ""); 65 | 66 | [[nodiscard]] bool partialEqual( 67 | const ConnectionKey& rhs) const noexcept override; 68 | 69 | [[nodiscard]] bool operator==(const MysqlConnectionKey& rhs) const noexcept; 70 | 71 | [[nodiscard]] bool operator!=(const MysqlConnectionKey& rhs) const noexcept { 72 | return !(*this == rhs); 73 | } 74 | 75 | [[nodiscard]] virtual bool operator==( 76 | const ConnectionKey& rhs) const noexcept override; 77 | 78 | [[nodiscard]] const std::string& host() const override { 79 | return host_; 80 | } 81 | 82 | [[nodiscard]] const std::string& db_name() const noexcept override { 83 | return dbName_; 84 | } 85 | 86 | [[nodiscard]] const std::string& user() const noexcept override { 87 | return user_; 88 | } 89 | 90 | [[nodiscard]] const std::string& password() const noexcept override { 91 | return password_; 92 | } 93 | 94 | [[nodiscard]] const std::string& unixSocketPath() const noexcept { 95 | return unixSocketPath_; 96 | } 97 | 98 | [[nodiscard]] size_t hash() const noexcept override { 99 | return hash_; 100 | } 101 | 102 | [[nodiscard]] size_t partial_hash() const noexcept override { 103 | return partialHash_; 104 | } 105 | 106 | [[nodiscard]] int port() const override { 107 | return port_; 108 | } 109 | 110 | [[nodiscard]] const std::string& special_tag() const noexcept override { 111 | return specialTag_; 112 | } 113 | 114 | [[nodiscard]] std::string getDisplayString( 115 | bool level2 = false) const override; 116 | 117 | private: 118 | std::string host_; 119 | std::string dbName_; 120 | std::string user_; 121 | std::string password_; 122 | std::string specialTag_; 123 | std::string unixSocketPath_; 124 | size_t partialHash_; 125 | size_t hash_; 126 | int port_; 127 | bool ignoreDbName_; 128 | }; 129 | 130 | } // namespace facebook::common::mysql_client 131 | 132 | // make default template of unordered_map/unordered_set works for ConnectionKey 133 | namespace std { 134 | template <> 135 | struct hash { 136 | size_t operator()( 137 | const facebook::common::mysql_client::ConnectionKey& k) const { 138 | return k.hash(); 139 | } 140 | }; 141 | 142 | template <> 143 | struct hash { 144 | size_t operator()( 145 | const facebook::common::mysql_client::MysqlConnectionKey& k) const { 146 | return k.hash(); 147 | } 148 | }; 149 | } // namespace std 150 | -------------------------------------------------------------------------------- /squangle/base/ExceptionUtil.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | 14 | namespace facebook::db { 15 | 16 | class FOLLY_EXPORT Exception : public std::runtime_error { 17 | public: 18 | using std::runtime_error::runtime_error; 19 | }; 20 | 21 | class FOLLY_EXPORT OperationStateException : public Exception { 22 | public: 23 | using Exception::Exception; 24 | }; 25 | 26 | class FOLLY_EXPORT InvalidConnectionException : public Exception { 27 | public: 28 | using Exception::Exception; 29 | }; 30 | 31 | } // namespace facebook::db 32 | -------------------------------------------------------------------------------- /squangle/logger/DBEventCounter.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/logger/DBEventCounter.h" 10 | 11 | #include 12 | #include 13 | #include 14 | 15 | namespace facebook { 16 | namespace db { 17 | 18 | void ConnectionContextBase::collectValues( 19 | const AddIntValueFunction& addInt, 20 | const AddNormalValueFunction& addNormal) const { 21 | addInt("ssl_server_cert_validated", isServerCertValidated ? 1 : 0); 22 | addInt("ssl_client_identity_cert", isIdentityClientCert ? 1 : 0); 23 | if (certCacheSize.has_value()) { 24 | addInt("ssl_cert_cache_size", certCacheSize.value()); 25 | } 26 | addNormal("is_ssl", folly::to(isSslConnection)); 27 | addNormal("is_ssl_session_reused", folly::to(sslSessionReused)); 28 | if (!sslVersion.empty()) { 29 | addNormal("ssl_version", sslVersion); 30 | } 31 | if (sslCertCn.hasValue()) { 32 | addNormal("ssl_server_cert_cn", sslCertCn.value()); 33 | } 34 | if (sslCertSan.hasValue() && !sslCertSan.value().empty()) { 35 | addNormal("ssl_server_cert_san", folly::join(',', sslCertSan.value())); 36 | } 37 | if (sslCertIdentities.hasValue() && !sslCertIdentities.value().empty()) { 38 | addNormal( 39 | "ssl_server_cert_identities", 40 | folly::join(',', sslCertIdentities.value())); 41 | } 42 | if (!endpointVersion.empty()) { 43 | addNormal("endpoint_version", endpointVersion); 44 | } 45 | } 46 | 47 | folly::Optional ConnectionContextBase::getNormalValue( 48 | folly::StringPiece key) const { 49 | if (key == "is_ssl") { 50 | return folly::to(isSslConnection); 51 | } else if (key == "is_ssl_session_reused") { 52 | return folly::to(sslSessionReused); 53 | } else if (key == "ssl_version" && !sslVersion.empty()) { 54 | return sslVersion; 55 | } else if (key == "ssl_server_cert_cn") { 56 | return sslCertCn; 57 | } else if (key == "ssl_server_cert_san") { 58 | if (sslCertSan.hasValue()) { 59 | return folly::join(',', sslCertSan.value()); 60 | } else { 61 | return folly::none; 62 | } 63 | } else if (key == "ssl_server_cert_identities") { 64 | if (sslCertIdentities.hasValue()) { 65 | return folly::join(',', sslCertIdentities.value()); 66 | } else { 67 | return folly::none; 68 | } 69 | } else if (key == "endpoint_version" && !endpointVersion.empty()) { 70 | return endpointVersion; 71 | } else if (key == "ssl_client_identity_cert") { 72 | return folly::to(isIdentityClientCert); 73 | } else if (key == "ssl_cert_cache_size" && certCacheSize.has_value()) { 74 | return folly::to(certCacheSize.value()); 75 | } else { 76 | return folly::none; 77 | } 78 | } 79 | 80 | ExponentialMovingAverage::ExponentialMovingAverage(double smoothingFactor) 81 | : smoothingFactor_(smoothingFactor) {} 82 | 83 | void ExponentialMovingAverage::addSample(double sample) { 84 | if (hasRegisteredFirstSample_) { 85 | currentValue_ = 86 | smoothingFactor_ * sample + (1 - smoothingFactor_) * currentValue_; 87 | } else { 88 | currentValue_ = sample; 89 | hasRegisteredFirstSample_ = true; 90 | } 91 | } 92 | 93 | void SimpleDbCounter::printStats() { 94 | LOG(INFO) << "Client Stats\n" 95 | << "Opened Connections " << numOpenedConnections() << "\n" 96 | << "Closed Connections " << numClosedConnections() << "\n" 97 | << "Failed Queries " << numFailedQueries() << "\n" 98 | << "Succeeded Queries " << numSucceededQueries() << "\n" 99 | << "Reused SSL Sessions " << numReusedSSLSessions() << "\n"; 100 | } 101 | } // namespace db 102 | } // namespace facebook 103 | -------------------------------------------------------------------------------- /squangle/logger/DBEventLogger.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/logger/DBEventLogger.h" 10 | 11 | #include 12 | 13 | namespace facebook { 14 | namespace db { 15 | 16 | void DBSimpleLogger::logQuerySuccess( 17 | const QueryLoggingData& data, 18 | const SquangleLoggingData&) const { 19 | VLOG(2) << "[" << api_name_ << "]" << " query (\"" << data.query 20 | << "\") succeeded."; 21 | } 22 | 23 | void DBSimpleLogger::logQueryFailure( 24 | const QueryLoggingData& data, 25 | FailureReason, 26 | unsigned int, 27 | const std::string&, 28 | const SquangleLoggingData&) const { 29 | VLOG(2) << "[" << api_name_ << "]" << " query (\"" << data.query 30 | << "\") failed."; 31 | } 32 | 33 | void DBSimpleLogger::logConnectionSuccess( 34 | const CommonLoggingData&, 35 | const SquangleLoggingData& connInfo) const { 36 | VLOG(2) << "[" << api_name_ << "]" << " connection with " 37 | << connInfo.getConnKeyRef().host() << " succeeded"; 38 | } 39 | 40 | void DBSimpleLogger::logConnectionFailure( 41 | const CommonLoggingData&, 42 | FailureReason, 43 | unsigned int, 44 | const std::string&, 45 | const SquangleLoggingData& connInfo) const { 46 | VLOG(2) << "[" << api_name_ << "]" << " connection with " 47 | << connInfo.getConnKeyRef().host() << " failed"; 48 | } 49 | } // namespace db 50 | } // namespace facebook 51 | -------------------------------------------------------------------------------- /squangle/mysql_client/AsyncConnectionPool.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | #include "squangle/mysql_client/AsyncConnectionPool.h" 18 | #include "squangle/mysql_client/AsyncMysqlClient.h" 19 | #include "squangle/mysql_client/ConnectPoolOperation.h" 20 | #include "squangle/mysql_client/SemiFutureAdapter.h" 21 | 22 | namespace facebook::common::mysql_client { 23 | 24 | std::shared_ptr AsyncConnectionPool::makePool( 25 | std::shared_ptr mysql_client, 26 | PoolOptions pool_options) { 27 | auto connectionPool = std::make_shared( 28 | std::move(mysql_client), std::move(pool_options)); 29 | return connectionPool; 30 | } 31 | 32 | AsyncConnectionPool::AsyncConnectionPool( 33 | std::shared_ptr mysql_client, 34 | PoolOptions pool_options) 35 | : ConnectionPool( 36 | std::move(mysql_client), 37 | std::move(pool_options)), 38 | cleanup_timer_(mysql_client_->getEventBase(), conn_storage_) { 39 | if (!mysql_client_->runInThread([this]() { 40 | cleanup_timer_.scheduleTimeout(PoolOptions::kCleanUpTimeout); 41 | })) { 42 | LOG(DFATAL) << "Unable to schedule timeout due Thrift event issue"; 43 | } 44 | } 45 | 46 | AsyncConnectionPool::~AsyncConnectionPool() { 47 | VLOG(2) << "Connection pool dying"; 48 | 49 | if (!shutdown_data_.rlock()->finished_shutdown) { 50 | shutdown(); 51 | } 52 | 53 | VLOG(2) << "Connection pool shutdown completed"; 54 | } 55 | 56 | void AsyncConnectionPool::shutdown() { 57 | VLOG(2) << "Shutting down"; 58 | 59 | auto shutdown_func = [&](auto& shutdown_data) { 60 | cleanup_timer_.cancelTimeout(); 61 | conn_storage_.clearAll(); 62 | shutdown_data.finished_shutdown = true; 63 | VLOG(1) << "Shutting down in mysql_client thread"; 64 | }; 65 | 66 | // New scope to limit the lifetime of the write lock 67 | { 68 | auto shutdown_data = shutdown_data_.wlock(); 69 | if (shutdown_data->shutting_down) { 70 | return; 71 | } 72 | 73 | // Will block adding anything to the pool 74 | shutdown_data->shutting_down = true; 75 | 76 | // cancelTimeout can only be ran in the mysql_client thread 77 | if (std::this_thread::get_id() == mysql_client_->threadId()) { 78 | shutdown_func(*shutdown_data); 79 | return; 80 | } 81 | } 82 | 83 | // We aren't already in the right thread, cause the shutdown to get run in the 84 | // correct thread and wait for it to complete. 85 | mysql_client_->runInThread( 86 | [&]() { shutdown_func(*shutdown_data_.wlock()); }, /*wait*/ true); 87 | } 88 | 89 | folly::SemiFuture AsyncConnectionPool::connectSemiFuture( 90 | const std::string& host, 91 | int port, 92 | const std::string& database_name, 93 | const std::string& user, 94 | const std::string& password, 95 | const ConnectionOptions& conn_opts) { 96 | return connectSemiFuture( 97 | host, port, database_name, user, password, "", conn_opts); 98 | } 99 | 100 | folly::SemiFuture AsyncConnectionPool::connectSemiFuture( 101 | const std::string& host, 102 | int port, 103 | const std::string& database_name, 104 | const std::string& user, 105 | const std::string& password, 106 | const std::string& special_tag, 107 | const ConnectionOptions& conn_opts) { 108 | auto op = 109 | beginConnection(host, port, database_name, user, password, special_tag); 110 | op->setConnectionOptions(conn_opts); 111 | return toSemiFuture(std::move(op)); 112 | } 113 | 114 | std::unique_ptr AsyncConnectionPool::connect( 115 | const std::string& host, 116 | int port, 117 | const std::string& database_name, 118 | const std::string& user, 119 | const std::string& password, 120 | const ConnectionOptions& conn_opts) { 121 | auto op = beginConnection(host, port, database_name, user, password); 122 | op->setConnectionOptions(conn_opts); 123 | // This will throw (intended behaviour) in case the operation didn't succeed 124 | return blockingConnectHelper(*op); 125 | } 126 | 127 | template <> 128 | void AsyncConnectPoolOperationImpl::specializedRun() { 129 | std::weak_ptr weakSelf = getOp().getSharedPointer(); 130 | if (!conn().client().runInThread([&, wself = std::move(weakSelf)]() { 131 | // There is a race confition that allows a cancelled or completed 132 | // operation getting here. The self ptr check ensures that the client 133 | // has not freed the reference to the operation, and the state() check 134 | // verifies whether other relevant memebers have been cleaned up by 135 | // connect callbacks 136 | auto self = wself.lock(); 137 | if (!self || (self->state() == OperationState::Completed)) { 138 | LOG(ERROR) << "ConnectPoolOperation freed before running"; 139 | return; 140 | } 141 | 142 | // `this` is valid if `self` is still alive 143 | specializedRunImpl(); 144 | })) { 145 | completeOperationInner(OperationResult::Failed); 146 | } 147 | } 148 | 149 | std::ostream& operator<<(std::ostream& os, ExpirationPolicy policy) { 150 | auto str = (policy == ExpirationPolicy::Age) 151 | ? "Age" 152 | : (policy == ExpirationPolicy::IdleTime ? "IdleTime" : ""); 153 | return os << str; 154 | } 155 | 156 | std::ostream& operator<<(std::ostream& os, const PoolOptions& options) { 157 | return os << "{per key limit:" << options.getPerKeyLimit() 158 | << ",pool limit:" << options.getPoolLimit() 159 | << ",idle timeout:" << options.getIdleTimeout().count() 160 | << "us,age timeout:" << options.getAgeTimeout().count() 161 | << "us,expiration policy:" << options.getExpPolicy() 162 | << ",pool per instance:" << options.poolPerMysqlInstance() << "}"; 163 | } 164 | 165 | std::ostream& operator<<(std::ostream& os, const PoolKey& key) { 166 | return os << "{key:" << key.getConnectionKeyRef().getDisplayString() 167 | << ",options:" << key.getConnectionOptions().getDisplayString() 168 | << "}"; 169 | } 170 | 171 | template <> 172 | std::unique_ptr> 173 | createConnectPoolOperationImpl( 174 | std::weak_ptr> pool, 175 | std::shared_ptr client, 176 | std::shared_ptr conn_key) { 177 | return std::make_unique< 178 | mysql_protocol::MysqlConnectPoolOperationImpl>( 179 | std::move(pool), client, std::move(conn_key)); 180 | } 181 | 182 | } // namespace facebook::common::mysql_client 183 | -------------------------------------------------------------------------------- /squangle/mysql_client/AsyncConnectionPool.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | // 10 | // Asynchronous Connection Pool based on our async MySQL client. 11 | // 12 | // This pool offers na async way to acquire a connection by creating new ones 13 | // or recycling an existing one. Also provides a way to limit the number of 14 | // open connections per database/user and for the client. 15 | // 16 | // AsyncConnectionPool - This pool holds multiple MySQL connections and 17 | // manages them to make sure only healthy connections are given back. 18 | // The interface to request a connection works just like the 19 | // AsyncMysqlClient, an ConnectPoolOperation is started by `beginConnection`. 20 | // 21 | // ConnectPoolOperation - An abstraction of ConnectOperation that instead of 22 | // opening a new connection, requests a connection to the pool it was created 23 | // by. The usage and error treat are the same. 24 | 25 | #pragma once 26 | 27 | #include 28 | #include 29 | #include 30 | #include 31 | #include 32 | 33 | #include "squangle/mysql_client/AsyncMysqlClient.h" 34 | #include "squangle/mysql_client/ConnectPoolOperation.h" 35 | #include "squangle/mysql_client/ConnectionPool.h" 36 | #include "squangle/mysql_client/Operation.h" 37 | #include "squangle/mysql_client/mysql_protocol/MysqlConnectPoolOperationImpl.h" 38 | 39 | namespace facebook::common::mysql_client { 40 | 41 | template 42 | class ConnectPoolOperation; 43 | class AsyncConnectionPool; 44 | class PoolKey; 45 | 46 | using AsyncConnectPoolOperation = ConnectPoolOperation; 47 | using AsyncConnectPoolOperationImpl = 48 | mysql_protocol::MysqlConnectPoolOperationImpl; 49 | 50 | class AsyncConnectionPool : public ConnectionPool { 51 | public: 52 | // Don't use std::chrono::duration::MAX to avoid overflows 53 | static std::shared_ptr makePool( 54 | std::shared_ptr mysql_client, 55 | PoolOptions pool_options = PoolOptions()); 56 | 57 | // The destructor will start the shutdown phase 58 | ~AsyncConnectionPool() override; 59 | 60 | // Don't allow copy or move 61 | AsyncConnectionPool(const AsyncConnectionPool&) = delete; 62 | AsyncConnectionPool& operator=(const AsyncConnectionPool&) = delete; 63 | 64 | AsyncConnectionPool(AsyncConnectionPool&&) = delete; 65 | AsyncConnectionPool& operator=(AsyncConnectionPool&&) = delete; 66 | 67 | FOLLY_NODISCARD folly::SemiFuture connectSemiFuture( 68 | const std::string& host, 69 | int port, 70 | const std::string& database_name, 71 | const std::string& user, 72 | const std::string& password, 73 | const ConnectionOptions& conn_opts = ConnectionOptions()); 74 | 75 | FOLLY_NODISCARD folly::SemiFuture connectSemiFuture( 76 | const std::string& host, 77 | int port, 78 | const std::string& database_name, 79 | const std::string& user, 80 | const std::string& password, 81 | const std::string& special_tag, 82 | const ConnectionOptions& conn_opts = ConnectionOptions()); 83 | 84 | std::unique_ptr connect( 85 | const std::string& host, 86 | int port, 87 | const std::string& database_name, 88 | const std::string& user, 89 | const std::string& password, 90 | const ConnectionOptions& conn_opts = ConnectionOptions()); 91 | 92 | // Don't use the constructor directly, only public to use make_shared 93 | AsyncConnectionPool( 94 | std::shared_ptr mysql_client, 95 | PoolOptions pool_options); 96 | 97 | // It will clean the pool and block any new connections or operations 98 | // Shutting down phase: 99 | // The remaining connections or operations that are linked to this pool 100 | // will know (using their weak_pointer to this pool) that the pool is dead 101 | // and proceed without the pool. 102 | void shutdown() override; 103 | 104 | // for debugging, return number of pool keys in the pool 105 | FOLLY_NODISCARD size_t getNumKey() const noexcept { 106 | return conn_storage_.getNumKey(); 107 | } 108 | 109 | private: 110 | template 111 | friend class ConnectPoolOperation; 112 | 113 | class CleanUpTimer : public folly::AsyncTimeout { 114 | public: 115 | explicit CleanUpTimer( 116 | folly::EventBase* base, 117 | PoolStorage& pool) 118 | : folly::AsyncTimeout(base), pool_(pool) {} 119 | 120 | void timeoutExpired() noexcept override { 121 | scheduleTimeout(PoolOptions::kCleanUpTimeout); 122 | pool_.cleanupOperations(); 123 | pool_.cleanupConnections(); 124 | } 125 | 126 | private: 127 | PoolStorage& pool_; 128 | } cleanup_timer_; 129 | 130 | // The AsyncConnectionPool needs to make sure certain things run in the mysql 131 | // thread. 132 | void validateCorrectThread() const override { 133 | DCHECK_EQ(std::this_thread::get_id(), mysql_client_->threadId()); 134 | } 135 | 136 | bool runInCorrectThread(std::function&& func, bool /*wait*/) 137 | override { 138 | return mysql_client_->runInThread([func = std::move(func)]() { func(); }); 139 | } 140 | 141 | std::unique_ptr makeNewConnection( 142 | std::shared_ptr conn_key, 143 | std::unique_ptr> mysqlConn) override { 144 | return std::make_unique( 145 | *mysql_client_, std::move(conn_key), std::move(mysqlConn)); 146 | } 147 | 148 | struct ShutdownData { 149 | bool shutting_down = false; 150 | bool finished_shutdown = false; 151 | }; 152 | 153 | folly::Synchronized shutdown_data_; 154 | 155 | bool isShuttingDown() const override { 156 | return shutdown_data_.rlock()->shutting_down; 157 | } 158 | 159 | // Nothing needed here - the async connection pool should not wait. 160 | void openNewConnectionPrep(AsyncConnectPoolOperation& /*pool_op*/) override {} 161 | void openNewConnectionFinish( 162 | AsyncConnectPoolOperation& /*pool_op*/, 163 | const PoolKey& /*pool_key*/) override {} 164 | }; 165 | 166 | } // namespace facebook::common::mysql_client 167 | -------------------------------------------------------------------------------- /squangle/mysql_client/AsyncHelpers.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/AsyncHelpers.h" 10 | 11 | namespace facebook::common::mysql_client { 12 | 13 | QueryCallback resultAppender(QueryAppenderCallback&& callback) { 14 | return [callback = std::move(callback)]( 15 | QueryOperation& op, 16 | QueryResult* /*res*/, 17 | QueryCallbackReason reason) { 18 | if (reason != QueryCallbackReason::RowsFetched) { 19 | QueryResult result{0}; 20 | if (op.ok()) { 21 | // Can't access query result on error 22 | result = op.stealQueryResult(); 23 | } 24 | callback(op, std::move(result), reason); 25 | } 26 | }; 27 | } 28 | 29 | MultiQueryCallback resultAppender(MultiQueryAppenderCallback&& callback) { 30 | return [callback = std::move(callback)]( 31 | MultiQueryOperation& op, 32 | QueryResult* /*res*/, 33 | QueryCallbackReason reason) { 34 | if (reason != QueryCallbackReason::RowsFetched && 35 | reason != QueryCallbackReason::QueryBoundary) { 36 | // stealQueryResults is always allowed since in multiQuery 37 | // the first few queries might succeed and then be followed by a failure 38 | callback(op, op.stealQueryResults(), reason); 39 | } 40 | }; 41 | } 42 | 43 | } // namespace facebook::common::mysql_client 44 | -------------------------------------------------------------------------------- /squangle/mysql_client/AsyncHelpers.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | 14 | #include "squangle/mysql_client/MultiQueryOperation.h" 15 | #include "squangle/mysql_client/QueryOperation.h" 16 | 17 | namespace facebook::common::mysql_client { 18 | 19 | // This will allow callback users to be able to read all rows in the end of 20 | // operation (Success or Failure). 21 | // The Appender callbacks will only be called for `Success` and `Failure` 22 | // and pass the whole result of the query or multi query. 23 | using QueryAppenderCallback = 24 | std::function; 25 | using MultiQueryAppenderCallback = std::function< 26 | void(MultiQueryOperation&, std::vector, QueryCallbackReason)>; 27 | 28 | QueryCallback resultAppender(QueryAppenderCallback&& callback); 29 | 30 | MultiQueryCallback resultAppender(MultiQueryAppenderCallback&& callback); 31 | 32 | } // namespace facebook::common::mysql_client 33 | -------------------------------------------------------------------------------- /squangle/mysql_client/ChangeUserOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/ChangeUserOperation.h" 10 | #include "squangle/mysql_client/Connection.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | InternalConnection::Status ChangeUserOperation::runSpecialOperation() { 15 | return conn().getInternalConnection().changeUser(key_); 16 | } 17 | 18 | } // namespace facebook::common::mysql_client 19 | -------------------------------------------------------------------------------- /squangle/mysql_client/ChangeUserOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/SpecialOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class ChangeUserOperation : public SpecialOperation { 16 | public: 17 | explicit ChangeUserOperation( 18 | std::unique_ptr impl, 19 | std::shared_ptr key) 20 | : SpecialOperation(std::move(impl)), key_(std::move(key)) {} 21 | 22 | private: 23 | InternalConnection::Status runSpecialOperation() override; 24 | 25 | db::OperationType getOperationType() const override { 26 | return db::OperationType::ChangeUser; 27 | } 28 | 29 | const char* getErrorMsg() const override { 30 | return errorMsg; 31 | } 32 | 33 | std::shared_ptr key_; 34 | 35 | static constexpr const char* errorMsg = "Change user failed: "; 36 | }; 37 | 38 | } // namespace facebook::common::mysql_client 39 | -------------------------------------------------------------------------------- /squangle/mysql_client/ClientPool.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | 15 | namespace facebook::common::mysql_client { 16 | 17 | class AsyncMysqlClientFactory { 18 | public: 19 | std::shared_ptr makeClient() { 20 | return std::shared_ptr( 21 | new AsyncMysqlClient(), AsyncMysqlClient::deleter); 22 | } 23 | }; 24 | 25 | /* 26 | * For MySQL heavy applications, we require more AsyncMysqlClient's. 27 | * Given that each spins a thread. 28 | * This is a simple round robin of pools. 29 | */ 30 | // TClientFactory is here just to allow the PoolOptions to be passed as well. 31 | template 32 | class ClientPool { 33 | public: 34 | explicit ClientPool( 35 | std::unique_ptr client_factory, 36 | size_t num_clients = 10) { 37 | if (num_clients == 0) { 38 | throw std::logic_error( 39 | "Invalid number of clients, it needs to be more than 0"); 40 | } 41 | client_pool_.reserve(num_clients); 42 | for (int i = 0; i < num_clients; ++i) { 43 | client_pool_.emplace_back(client_factory->makeClient()); 44 | } 45 | } 46 | 47 | std::shared_ptr getClient() const { 48 | auto idx = folly::Random::rand32() % client_pool_.size(); 49 | return client_pool_[idx]; 50 | } 51 | 52 | template 53 | void forEachClient(F func) { 54 | for (auto& client : client_pool_) { 55 | func(client); 56 | } 57 | } 58 | 59 | // Passing in a key will allow the use of a consistent AsyncConnectionPool 60 | // object per key. This will greatly increase pool hits as currently 61 | // the multiple pools do not share any resources. This also allows the 62 | // MultiPool to respect limits 63 | std::shared_ptr getClient(const std::string& key) const { 64 | return getClient(folly::Hash()(key)); 65 | } 66 | 67 | // Using size_t key to index the client pool 68 | std::shared_ptr getClient(size_t key) const { 69 | return client_pool_[key % client_pool_.size()]; 70 | } 71 | 72 | static std::shared_ptr getClientFromDefault() { 73 | auto client_pool = 74 | folly::Singleton>::try_get(); 75 | if (!client_pool) { 76 | throw std::logic_error( 77 | "MultiMysqlClientPool singleton has already been destroyed."); 78 | } 79 | return client_pool->getClient(); 80 | } 81 | 82 | static std::shared_ptr getClientFromDefault(const std::string& key) { 83 | auto client_pool = 84 | folly::Singleton>::try_get(); 85 | if (!client_pool) { 86 | throw std::logic_error( 87 | "MultiMysqlClientPool singleton has already been destroyed."); 88 | } 89 | return client_pool->getClient(key); 90 | } 91 | 92 | private: 93 | std::vector> client_pool_; 94 | }; 95 | 96 | } // namespace facebook::common::mysql_client 97 | -------------------------------------------------------------------------------- /squangle/mysql_client/Compression.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include "squangle/mysql_client/Compression.h" 15 | 16 | namespace facebook::common::mysql_client { 17 | 18 | bool setCompressionOption(MYSQL* mysql, CompressionAlgorithm algo) { 19 | #if MYSQL_VERSION_ID < 80018 20 | auto comp_lib = getCompressionValue(algo); 21 | auto res = mysql_options(mysql, MYSQL_OPT_COMP_LIB, (void*)&comp_lib); 22 | #else 23 | auto res = mysql_options( 24 | mysql, MYSQL_OPT_COMPRESSION_ALGORITHMS, getCompressionName(algo).data()); 25 | #endif 26 | return res == 0; // 0 indicates success 27 | } 28 | 29 | #if MYSQL_VERSION_ID < 80018 30 | static const std::unordered_map 31 | compressionAlgorithms = { 32 | {ZLIB, MYSQL_COMPRESSION_ZLIB}, 33 | {ZSTD, MYSQL_COMPRESSION_ZSTD}, 34 | {ZSTD_STREAM, MYSQL_COMPRESSION_ZSTD_STREAM}, 35 | {LZ4F_STREAM, MYSQL_COMPRESSION_LZ4F_STREAM}, 36 | }; 37 | 38 | mysql_compression_lib getCompressionValue(CompressionAlgorithm algo) { 39 | auto it = compressionAlgorithms.find(algo); 40 | CHECK(it != compressionAlgorithms.end()) 41 | << "Invalid compression algorithm enum: " << (int)algo; 42 | 43 | return it->second; 44 | } 45 | #else 46 | // These strings _must_ match the string supported in MySQL 47 | // (https://fburl.com/diffusion/5jkzvbg0) 48 | static const std::unordered_map 49 | compressionAlgorithms = { 50 | {ZLIB, "zlib"}, 51 | {ZSTD, "zstd"}, 52 | {ZSTD_STREAM, "zstd_stream"}, 53 | {LZ4F_STREAM, "lz4f_stream"}, 54 | }; 55 | 56 | std::optional parseCompressionName( 57 | std::string_view name) { 58 | auto it = std::find_if( 59 | compressionAlgorithms.begin(), 60 | compressionAlgorithms.end(), 61 | [&](const auto& entry) { return entry.second == name; }); 62 | 63 | if (it == compressionAlgorithms.end()) { 64 | return std::nullopt; 65 | } 66 | 67 | return it->first; 68 | } 69 | 70 | const std::string& getCompressionName(CompressionAlgorithm algo) { 71 | auto it = compressionAlgorithms.find(algo); 72 | CHECK(it != compressionAlgorithms.end()) 73 | << "Invalid compression algorithm enum: " << (int)algo; 74 | 75 | return it->second; 76 | } 77 | #endif 78 | 79 | } // namespace facebook::common::mysql_client 80 | -------------------------------------------------------------------------------- /squangle/mysql_client/Compression.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include // @manual 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | // Compression algorithms supported by MySQL: 16 | // https://fburl.com/diffusion/kq2fbmr2 17 | enum CompressionAlgorithm { 18 | ZLIB, 19 | ZSTD, 20 | ZSTD_STREAM, 21 | LZ4F_STREAM, 22 | }; 23 | 24 | bool setCompressionOption(MYSQL* mysql, CompressionAlgorithm algo); 25 | 26 | #if MYSQL_VERSION_ID >= 80000 27 | #if MYSQL_VERSION_ID < 80018 28 | mysql_compression_lib getCompressionValue(CompressionAlgorithm algo); 29 | #else 30 | std::optional parseCompressionName(std::string_view name); 31 | const std::string& getCompressionName(CompressionAlgorithm algo); 32 | #endif 33 | #endif 34 | 35 | } // namespace facebook::common::mysql_client 36 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | #include "squangle/mysql_client/ConnectOperation.h" 13 | #include "squangle/mysql_client/Connection.h" 14 | 15 | namespace facebook::common::mysql_client { 16 | 17 | using namespace std::chrono_literals; 18 | 19 | ConnectOperationImpl::ConnectOperationImpl( 20 | MysqlClientBase* mysql_client, 21 | std::shared_ptr conn_key) 22 | : conn_key_(std::move(conn_key)) { 23 | mysql_client->activeConnectionAdded(conn_key_); 24 | } 25 | 26 | void ConnectOperationImpl::setConnectionOptions( 27 | const ConnectionOptions& conn_opts) { 28 | setTimeout(conn_opts.getTimeout()); 29 | setDefaultQueryTimeout(conn_opts.getQueryTimeout()); 30 | setAttributes(conn_opts.getAttributes()); 31 | setConnectAttempts(conn_opts.getConnectAttempts()); 32 | if (conn_opts.getDscp().has_value()) { 33 | setDscp(*conn_opts.getDscp()); 34 | } 35 | setTotalTimeout(conn_opts.getTotalTimeout()); 36 | setCompression(conn_opts.getCompression()); 37 | if (conn_opts.getConnectTcpTimeout()) { 38 | setTcpTimeout(*conn_opts.getConnectTcpTimeout()); 39 | } 40 | if (conn_opts.getSniServerName()) { 41 | setSniServerName(*conn_opts.getSniServerName()); 42 | } 43 | if (auto provider = conn_opts.getSSLOptionsProvider()) { 44 | setSSLOptionsProvider(std::move(provider)); 45 | } 46 | if (conn_opts.getCertValidationCallback()) { 47 | setCertValidationCallback(conn_opts.getCertValidationCallback()); 48 | } 49 | if (auto cats = conn_opts.getCryptoAuthTokenList()) { 50 | setCryptoAuthTokenList(*std::move(cats)); 51 | } 52 | } 53 | 54 | const ConnectionOptions& ConnectOperationImpl::getConnectionOptions() const { 55 | return conn_options_; 56 | } 57 | 58 | void ConnectOperationImpl::setDefaultQueryTimeout(Duration t) { 59 | CHECK_THROW( 60 | state() == OperationState::Unstarted, db::OperationStateException); 61 | conn_options_.setQueryTimeout(t); 62 | } 63 | 64 | void ConnectOperationImpl::setSniServerName(const std::string& sni_servername) { 65 | CHECK_THROW( 66 | state() == OperationState::Unstarted, db::OperationStateException); 67 | conn_options_.setSniServerName(sni_servername); 68 | } 69 | 70 | void ConnectOperationImpl::enableResetConnBeforeClose() { 71 | conn_options_.enableResetConnBeforeClose(); 72 | } 73 | 74 | void ConnectOperationImpl::enableDelayedResetConn() { 75 | conn_options_.enableDelayedResetConn(); 76 | } 77 | 78 | void ConnectOperationImpl::enableChangeUser() { 79 | conn_options_.enableChangeUser(); 80 | } 81 | 82 | void ConnectOperationImpl::setCertValidationCallback( 83 | CertValidatorCallback callback) { 84 | CHECK_THROW( 85 | state() == OperationState::Unstarted, db::OperationStateException); 86 | conn_options_.setCertValidationCallback(std::move(callback)); 87 | } 88 | 89 | void ConnectOperationImpl::setTimeout(Duration timeout) { 90 | conn_options_.setTimeout(timeout); 91 | OperationBase::setTimeout(timeout); 92 | } 93 | 94 | void ConnectOperationImpl::setTcpTimeout(Duration timeout) { 95 | conn_options_.setConnectTcpTimeout(timeout); 96 | } 97 | 98 | void ConnectOperationImpl::setTotalTimeout(Duration total_timeout) { 99 | conn_options_.setTotalTimeout(total_timeout); 100 | OperationBase::setTimeout(min(getTimeout(), total_timeout)); 101 | } 102 | void ConnectOperationImpl::setConnectAttempts(uint32_t max_attempts) { 103 | CHECK_THROW( 104 | state() == OperationState::Unstarted, db::OperationStateException); 105 | conn_options_.setConnectAttempts(max_attempts); 106 | } 107 | 108 | void ConnectOperationImpl::setDscp(uint8_t dscp) { 109 | CHECK_THROW( 110 | state() == OperationState::Unstarted, db::OperationStateException); 111 | conn_options_.setDscp(dscp); 112 | } 113 | 114 | void ConnectOperationImpl::setKillOnQueryTimeout(bool killOnQueryTimeout) { 115 | CHECK_THROW( 116 | state() == OperationState::Unstarted, db::OperationStateException); 117 | killOnQueryTimeout_ = killOnQueryTimeout; 118 | } 119 | void ConnectOperationImpl::setSSLOptionsProviderBase( 120 | std::unique_ptr /*ssl_options_provider*/) { 121 | CHECK_THROW( 122 | state() == OperationState::Unstarted, db::OperationStateException); 123 | LOG(ERROR) << "Using deprecated function"; 124 | } 125 | void ConnectOperationImpl::setSSLOptionsProvider( 126 | std::shared_ptr ssl_options_provider) { 127 | CHECK_THROW( 128 | state() == OperationState::Unstarted, db::OperationStateException); 129 | conn_options_.setSSLOptionsProvider(ssl_options_provider); 130 | } 131 | 132 | ConnectOperation& ConnectOperationImpl::op() const { 133 | DCHECK(op_ && dynamic_cast(op_) != nullptr); 134 | return *(ConnectOperation*)op_; 135 | } 136 | 137 | [[nodiscard]] const InternalConnection& 138 | ConnectOperationImpl::getInternalConnection() const { 139 | return conn().getInternalConnection(); 140 | } 141 | 142 | [[nodiscard]] InternalConnection& 143 | ConnectOperationImpl::getInternalConnection() { 144 | return conn().getInternalConnection(); 145 | } 146 | 147 | ConnectionHolder* ConnectOperationImpl::mysqlConnection() const { 148 | return conn().mysqlConnection(); 149 | } 150 | 151 | [[nodiscard]] std::string ConnectOperationImpl::generateTimeoutError( 152 | Millis elapsed, 153 | const std::function& getErrorNo, 154 | std::string apiName, 155 | std::optional location, 156 | std::optional extra) const { 157 | auto cbDelay = client_.callbackDelayAvg(); 158 | bool stalled = (cbDelay >= kCallbackDelayStallThreshold); 159 | 160 | // Overall the message looks like this: 161 | // [](Mysql Client) to : timed out [in pool] 162 | // [at stage ] (took Nms, timeout was Nms) 163 | // [(CLIENT_OVERLOADED: cb delay Nms, N active conns)] [TcpTimeout:N] 164 | std::vector parts; 165 | parts.push_back(fmt::format( 166 | "[{}]({}){} to {}:{} timed out", 167 | getErrorNo(stalled), 168 | kErrorPrefix, 169 | apiName, 170 | conn_key_->host(), 171 | conn_key_->port())); 172 | if (location) { 173 | parts.push_back(std::move(*location)); 174 | } 175 | parts.push_back(timeoutMessage(elapsed)); 176 | if (stalled) { 177 | parts.push_back(threadOverloadMessage(cbDelay)); 178 | } 179 | if (extra) { 180 | parts.push_back(std::move(*extra)); 181 | } 182 | 183 | return folly::join(" ", parts); 184 | } 185 | 186 | /*static*/ 187 | std::shared_ptr ConnectOperation::create( 188 | std::unique_ptr impl) { 189 | // We must do this unusual behavior (with `new`) instead of std::make_shared 190 | // because we don't want the constructor for ConnectOperation to be public. 191 | // Without a public constructor there is no standard way of allowing 192 | // std::make_shared to call the constructor - i.e. no way to mark 193 | // std::make_shared as a friend. So we have to do this weirdness. 194 | return std::shared_ptr( 195 | new ConnectOperation(std::move(impl))); 196 | } 197 | 198 | void ConnectOperation::callConnectCallback() { 199 | if (connect_callback_) { 200 | connect_callback_(*this); 201 | // Release callback since no other callbacks will be made 202 | connect_callback_ = nullptr; 203 | } 204 | } 205 | 206 | } // namespace facebook::common::mysql_client 207 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectPoolOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | #include "squangle/mysql_client/ConnectOperation.h" 14 | #include "squangle/mysql_client/PoolStorage.h" 15 | 16 | namespace facebook::common::mysql_client { 17 | 18 | template 19 | class ConnectionPool; 20 | 21 | template 22 | class PoolStorageData; 23 | 24 | template 25 | class MysqlPooledHolder; 26 | 27 | class AsyncConnectionPool; 28 | class SyncConnectionPool; 29 | class ConnectionHolder; 30 | struct PoolKeyStats; 31 | 32 | template 33 | class ConnectPoolOperationImpl : virtual public ConnectOperationImpl { 34 | public: 35 | virtual ~ConnectPoolOperationImpl() override = default; 36 | 37 | virtual void prepWait() = 0; 38 | virtual bool syncWait() = 0; 39 | virtual void cleanupWait() = 0; 40 | 41 | virtual void attemptFailed(OperationResult result) = 0; 42 | virtual void connectionCallback( 43 | std::unique_ptr> mysql_conn) = 0; 44 | }; 45 | 46 | template 47 | class ConnectPoolOperation : public ConnectOperation { 48 | public: 49 | explicit ConnectPoolOperation( 50 | std::unique_ptr> impl) 51 | : ConnectOperation(std::move(impl)) {} 52 | 53 | ~ConnectPoolOperation() override { 54 | cancelPreOperation(); 55 | } 56 | 57 | ConnectPoolOperation(const ConnectPoolOperation&) = delete; 58 | ConnectPoolOperation& operator=(const ConnectPoolOperation&) = delete; 59 | 60 | ConnectPoolOperation(ConnectPoolOperation&&) = delete; 61 | ConnectPoolOperation& operator=(ConnectPoolOperation&&) = delete; 62 | 63 | db::OperationType getOperationType() const override { 64 | return db::OperationType::PoolConnect; 65 | } 66 | 67 | void prepWait() { 68 | impl()->prepWait(); 69 | } 70 | 71 | bool syncWait() { 72 | return impl()->syncWait(); 73 | } 74 | 75 | void cleanupWait() { 76 | impl()->cleanupWait(); 77 | } 78 | 79 | bool setPreOperation(std::shared_ptr op) { 80 | return preOperation_.wlock()->set(std::move(op)); 81 | } 82 | 83 | void cancelPreOperation() { 84 | preOperation_.wlock()->cancel(); 85 | } 86 | 87 | void resetPreOperation() { 88 | preOperation_.wlock()->reset(); 89 | } 90 | 91 | protected: 92 | ConnectPoolOperationImpl* impl() override { 93 | return dynamic_cast*>( 94 | ConnectOperation::impl()); 95 | } 96 | const ConnectPoolOperationImpl* impl() const override { 97 | return dynamic_cast*>( 98 | ConnectOperation::impl()); 99 | } 100 | 101 | void attemptFailed(OperationResult result) { 102 | return impl()->attemptFailed(result); 103 | } 104 | 105 | private: 106 | friend class ConnectionPool; 107 | friend class PoolStorageData; 108 | 109 | // PreOperation keeps any other operation that needs to be canceled when 110 | // ConnectPoolOperation is cancelled. 111 | // PreOperation is not reused and its lifetime is with ConnectPoolOperation. 112 | // PreOperation is not thread-safe, and ConnectPoolOperation is responsible 113 | // for adding a lock. 114 | class PreOperation { 115 | public: 116 | void cancel() { 117 | if (preOperation_) { 118 | preOperation_->cancel(); 119 | preOperation_.reset(); 120 | } 121 | canceled_ = true; 122 | } 123 | 124 | void reset() { 125 | preOperation_.reset(); 126 | } 127 | 128 | // returns true if set pre-operation succeeds, false otherwise 129 | bool set(std::shared_ptr op) { 130 | if (canceled_) { 131 | return false; 132 | } 133 | preOperation_ = std::move(op); 134 | return true; 135 | } 136 | 137 | private: 138 | std::shared_ptr preOperation_; 139 | bool canceled_{false}; 140 | }; 141 | 142 | // Operation that is required before completing this operation, which could 143 | // be reset_connection or change_user operation. There's at most 1 144 | // pre-operation. 145 | folly::Synchronized preOperation_; 146 | 147 | void connectionCallback( 148 | std::unique_ptr> mysql_conn) { 149 | impl()->connectionCallback(std::move(mysql_conn)); 150 | } 151 | 152 | // Called when the connection that the pool is trying to acquire failed 153 | void failureCallback( 154 | OperationResult failure, 155 | unsigned int mysql_errno, 156 | const std::string& mysql_error) { 157 | setAsyncClientError(mysql_errno, mysql_error); 158 | attemptFailed(failure); 159 | } 160 | 161 | void specializedRunImpl() { 162 | impl()->specializedRunImpl(); 163 | } 164 | }; 165 | 166 | } // namespace facebook::common::mysql_client 167 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectionHolder.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/ConnectionHolder.h" 10 | #include "squangle/mysql_client/MysqlClientBase.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | ConnectionHolder::ConnectionHolder( 15 | MysqlClientBase& client, 16 | std::unique_ptr internalConn, 17 | std::shared_ptr key) 18 | : client_(client), 19 | internalConn_(std::move(internalConn)), 20 | key_(std::move(key)), 21 | opened_(false) { 22 | createTime_ = std::chrono::steady_clock::now(); 23 | client_.activeConnectionAdded(key_); 24 | } 25 | 26 | ConnectionHolder::ConnectionHolder( 27 | ConnectionHolder& other, 28 | std::shared_ptr key) 29 | : client_(other.client_), 30 | internalConn_(other.stealInternalConnection()), 31 | context_(other.context_), 32 | key_(std::move(key)), 33 | opened_(other.opened_), 34 | createTime_(other.createTime_), 35 | lastActiveTime_(other.lastActiveTime_) { 36 | client_.activeConnectionAdded(key_); 37 | } 38 | 39 | void ConnectionHolder::updateConnectionKey( 40 | std::shared_ptr key) { 41 | client_.activeConnectionRemoved(key_); 42 | key_ = std::move(key); 43 | client_.activeConnectionAdded(key_); 44 | } 45 | 46 | ConnectionHolder::~ConnectionHolder() { 47 | if (internalConn_) { 48 | if (auto func = internalConn_->getCloseFunction()) { 49 | if (!client_.runInThread([func = std::move(func)]() { func(); })) { 50 | LOG(DFATAL) 51 | << "Connection couldn't be closed: error in folly::EventBase"; 52 | } 53 | } 54 | 55 | onClose(); 56 | } 57 | client_.activeConnectionRemoved(key_); 58 | } 59 | 60 | void ConnectionHolder::onClose() { 61 | if (opened_) { 62 | client_.stats()->incrClosedConnections(context_.get()); 63 | } 64 | } 65 | 66 | void ConnectionHolder::connectionOpened() { 67 | opened_ = true; 68 | lastActiveTime_ = Clock::now(); 69 | 70 | client_.stats()->incrOpenedConnections(context_.get()); 71 | } 72 | 73 | } // namespace facebook::common::mysql_client 74 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectionOptions.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include "squangle/mysql_client/ConnectionOptions.h" 15 | #include "squangle/mysql_client/Flags.h" 16 | 17 | namespace facebook::common::mysql_client { 18 | 19 | class Operation; 20 | 21 | ConnectionOptions::ConnectionOptions() 22 | : connection_timeout_(FLAGS_async_mysql_connect_timeout_micros), 23 | total_timeout_(FLAGS_async_mysql_timeout_micros * 2), 24 | query_timeout_(FLAGS_async_mysql_timeout_micros) {} 25 | 26 | std::string ConnectionOptions::getDisplayString() const { 27 | // Reserve 4 + 2 extra elements 28 | folly::small_vector parts; 29 | 30 | parts.push_back( 31 | fmt::format("conn timeout={}us", connection_timeout_.count())); 32 | parts.push_back(fmt::format("query timeout={}us", query_timeout_.count())); 33 | parts.push_back(fmt::format("total timeout={}us", total_timeout_.count())); 34 | parts.push_back(fmt::format("conn attempts={}", max_attempts_)); 35 | if (dscp_.has_value()) { 36 | parts.push_back(fmt::format("outbound dscp={}", *dscp_)); 37 | } 38 | if (ssl_options_provider_ != nullptr) { 39 | parts.push_back(fmt::format( 40 | "SSL options provider={}", (void*)ssl_options_provider_.get())); 41 | } 42 | if (compression_lib_.has_value()) { 43 | parts.push_back(fmt::format( 44 | "compression library={}", (void*)compression_lib_.get_pointer())); 45 | } 46 | 47 | if (!attributes_.empty()) { 48 | std::vector substrings; 49 | for (const auto& [key, value] : attributes_) { 50 | substrings.push_back(fmt::format("{}={}", key, value)); 51 | } 52 | 53 | parts.push_back(fmt::format( 54 | "connection attributes=[{}]", folly::join(",", substrings))); 55 | } 56 | return fmt::format("({})", folly::join(", ", parts)); 57 | } 58 | 59 | // Sets the differentiated service code point (DSCP) on the underlying 60 | // connection, which has the effect of embedding it into outgoing packet ip 61 | // headers. The value may be used to classify and prioritize said traffic. 62 | // 63 | // Note: A DSCP value is 6 bits and is packed into an 8 bit field. Users must 64 | // specify the unpacked (unshifted) 6-bit value. 65 | // 66 | // Note: This implementation only supports IPv6. 67 | // 68 | // Also known as "Quality of Service" (QoS), "Type of Service", "Class of 69 | // Service" (COS). 70 | // 71 | // See also RFC 2474 [0] and RFC 3542 6.5 (IPv6 sockopt) [1] 72 | // [0]: https://tools.ietf.org/html/rfc2474 73 | // [1]: https://tools.ietf.org/html/rfc3542#section-6.5 74 | ConnectionOptions& ConnectionOptions::setDscp(uint8_t dscp) { 75 | CHECK_THROW((dscp & 0b11000000) == 0, std::invalid_argument); 76 | dscp_ = dscp; 77 | return *this; 78 | } 79 | 80 | ConnectionOptions& ConnectionOptions::setCertValidationCallback( 81 | CertValidatorCallback callback) noexcept { 82 | certValidationCallback_ = std::move(callback); 83 | return *this; 84 | } 85 | 86 | } // namespace facebook::common::mysql_client 87 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectionOptions.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | #include "squangle/base/Base.h" 14 | #include "squangle/mysql_client/Compression.h" 15 | 16 | namespace facebook::common::mysql_client { 17 | 18 | class Operation; 19 | class SSLOptionsProviderBase; 20 | 21 | using CertValidatorCallback = std::function& context, 24 | folly::StringPiece& errMsg)>; 25 | 26 | class ConnectionOptions { 27 | public: 28 | ConnectionOptions(); 29 | 30 | // Each attempt to acquire a connection will take at maximum this duration. 31 | // Use setTotalTimeout if you want to limit the timeout for all attempts. 32 | ConnectionOptions& setTimeout(Duration dur) noexcept { 33 | connection_timeout_ = dur; 34 | return *this; 35 | } 36 | 37 | Duration getTimeout() const noexcept { 38 | return connection_timeout_; 39 | } 40 | 41 | // This time out is used for each connect attempt and this is the maximum 42 | // time allowed for client<->server tcp handshake. This timeout allows client 43 | // to failfast in case the MYSQL server is not reachable, where we don't want 44 | // to wait for the entire connection_timeout_ which also includes time buffer 45 | // for ssl and MYSQL protocol exchange 46 | ConnectionOptions& setConnectTcpTimeout(Duration dur) noexcept { 47 | connection_tcp_timeout_ = dur; 48 | return *this; 49 | } 50 | 51 | const folly::Optional& getConnectTcpTimeout() const noexcept { 52 | return connection_tcp_timeout_; 53 | } 54 | 55 | // The connection created by these options will apply this query timeout 56 | // to all statements executed 57 | ConnectionOptions& setQueryTimeout(Duration dur) noexcept { 58 | query_timeout_ = dur; 59 | return *this; 60 | } 61 | 62 | Duration getQueryTimeout() const noexcept { 63 | return query_timeout_; 64 | } 65 | 66 | // Used to provide an SSLContext and SSL_Session provider 67 | ConnectionOptions& setSSLOptionsProvider( 68 | std::shared_ptr ssl_options_provider) noexcept { 69 | ssl_options_provider_ = std::move(ssl_options_provider); 70 | return *this; 71 | } 72 | 73 | std::shared_ptr getSSLOptionsProvider() 74 | const noexcept { 75 | return ssl_options_provider_; 76 | } 77 | 78 | template 79 | void withPossibleSSLOptionsProvider(Func&& func) { 80 | if (ssl_options_provider_) { 81 | func(*ssl_options_provider_); 82 | } 83 | } 84 | 85 | // Provide a Connection Attribute to be passed in the connection handshake 86 | ConnectionOptions& setAttribute(std::string_view attr, std::string value) { 87 | attributes_[attr] = std::move(value); 88 | return *this; 89 | } 90 | 91 | // MySQL 5.6 connection attributes. Sent at time of connect. 92 | const AttributeMap& getAttributes() const noexcept { 93 | return attributes_; 94 | } 95 | 96 | ConnectionOptions& setAttributes(const AttributeMap& attributes) { 97 | for (auto& [key, value] : attributes) { 98 | attributes_[key] = value; 99 | } 100 | return *this; 101 | } 102 | 103 | // Sorry for the weird API, there is no enum for compression = None 104 | ConnectionOptions& setCompression( 105 | folly::Optional comp_lib) noexcept { 106 | compression_lib_ = std::move(comp_lib); 107 | return *this; 108 | } 109 | 110 | const folly::Optional& getCompression() const noexcept { 111 | return compression_lib_; 112 | } 113 | 114 | ConnectionOptions& setUseChecksum(bool useChecksum) noexcept { 115 | use_checksum_ = useChecksum; 116 | return *this; 117 | } 118 | 119 | FOLLY_NODISCARD bool getUseChecksum() const noexcept { 120 | return use_checksum_; 121 | } 122 | 123 | // Sets the amount of attempts that will be tried in order to acquire the 124 | // connection. Each attempt will take at maximum the given timeout. To set 125 | // a global timeout that the operation shouldn't take more than, use 126 | // setTotalTimeout. 127 | // 128 | // This is no longer recommended for use, due to higher level retries 129 | ConnectionOptions& setConnectAttempts(uint32_t max_attempts) noexcept { 130 | max_attempts_ = max_attempts; 131 | return *this; 132 | } 133 | 134 | uint32_t getConnectAttempts() const noexcept { 135 | return max_attempts_; 136 | } 137 | 138 | ConnectionOptions& setDscp(uint8_t dscp); 139 | 140 | folly::Optional getDscp() const noexcept { 141 | return dscp_; 142 | } 143 | 144 | // If this is not set, but regular timeout was, the TotalTimeout for the 145 | // operation will be the number of attempts times the primary timeout. 146 | // Set this if you have strict timeout needs. 147 | // 148 | // This should generally not be set, as connectAttempts is 1 149 | ConnectionOptions& setTotalTimeout(Duration dur) noexcept { 150 | total_timeout_ = dur; 151 | return *this; 152 | } 153 | 154 | Duration getTotalTimeout() const noexcept { 155 | return total_timeout_; 156 | } 157 | 158 | std::string getDisplayString() const; 159 | 160 | ConnectionOptions& setSniServerName(const std::string& sniName) noexcept { 161 | sni_servername_ = sniName; 162 | return *this; 163 | } 164 | 165 | const folly::Optional& getSniServerName() const noexcept { 166 | return sni_servername_; 167 | } 168 | 169 | ConnectionOptions& enableResetConnBeforeClose() noexcept { 170 | reset_conn_before_close_ = true; 171 | return *this; 172 | } 173 | 174 | bool isEnableResetConnBeforeClose() const noexcept { 175 | return reset_conn_before_close_; 176 | } 177 | 178 | ConnectionOptions& enableDelayedResetConn() noexcept { 179 | delayed_reset_conn_ = true; 180 | return *this; 181 | } 182 | 183 | bool isEnableDelayedResetConn() const noexcept { 184 | return delayed_reset_conn_; 185 | } 186 | 187 | ConnectionOptions& enableChangeUser() noexcept { 188 | change_user_ = true; 189 | return *this; 190 | } 191 | 192 | bool isEnableChangeUser() const noexcept { 193 | return change_user_; 194 | } 195 | 196 | ConnectionOptions& setCertValidationCallback( 197 | CertValidatorCallback callback) noexcept; 198 | 199 | const CertValidatorCallback& getCertValidationCallback() const noexcept { 200 | return certValidationCallback_; 201 | } 202 | 203 | ConnectionOptions& setCryptoAuthTokenList(const std::string& tokenList) { 204 | crypt_auth_token_list_ = tokenList; 205 | return *this; 206 | } 207 | 208 | const folly::Optional getCryptoAuthTokenList() const { 209 | return crypt_auth_token_list_; 210 | } 211 | 212 | private: 213 | Duration connection_timeout_; 214 | folly::Optional connection_tcp_timeout_; 215 | Duration total_timeout_; 216 | Duration query_timeout_; 217 | std::shared_ptr ssl_options_provider_; 218 | AttributeMap attributes_; 219 | folly::Optional compression_lib_; 220 | bool use_checksum_ = false; 221 | uint32_t max_attempts_ = 1; 222 | folly::Optional dscp_; 223 | folly::Optional sni_servername_; 224 | folly::Optional crypt_auth_token_list_; 225 | bool reset_conn_before_close_ = false; 226 | bool delayed_reset_conn_ = false; 227 | bool change_user_ = false; 228 | CertValidatorCallback certValidationCallback_{nullptr}; 229 | }; 230 | 231 | } // namespace facebook::common::mysql_client 232 | -------------------------------------------------------------------------------- /squangle/mysql_client/ConnectionPool.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/ConnectionPool.h" 10 | 11 | namespace facebook::common::mysql_client { 12 | 13 | bool ConnectionPoolBase::canCreateMoreConnections( 14 | const PoolKey& pool_key, 15 | size_t enqueued_pool_ops, 16 | uint32_t client_total_conns, 17 | uint64_t client_conn_limit, 18 | const Counters& counters) const { 19 | // We have the number of connections we are opening and the number of already 20 | // open, we shouldn't try to create over this sum 21 | auto total_conns = 22 | counters.num_open_connections + counters.num_pending_connections; 23 | auto open_for_key = 24 | folly::get_default(counters.open_connections, pool_key, 0); 25 | auto pending_for_key = 26 | folly::get_default(counters.pending_connections, pool_key, 0); 27 | int conns_for_key = open_for_key + pending_for_key; 28 | 29 | // First we check global limit, then limits of the pool. If we can create more 30 | // connections, we check if we need comparing the amount of already being 31 | // opened connections for that key with the number of enqueued operations (the 32 | // operation that is requesting a new connection should be enqueued at this 33 | // point. 34 | return client_total_conns < client_conn_limit && total_conns < totalLimit() && 35 | conns_for_key < perKeyLimit() && pending_for_key < enqueued_pool_ops; 36 | } 37 | 38 | void ConnectionPoolBase::addOpenConnection(const PoolKey& pool_key) { 39 | counters_.withWLock([&](auto& locked) { 40 | ++locked.open_connections[pool_key]; 41 | ++locked.num_open_connections; 42 | }); 43 | } 44 | 45 | void ConnectionPoolBase::removeOpenConnection(const PoolKey& pool_key) { 46 | counters_.withWLock([&](auto& locked) { 47 | auto iter = locked.open_connections.find(pool_key); 48 | DCHECK(iter != locked.open_connections.end()); 49 | if (--iter->second == 0) { 50 | locked.open_connections.erase(iter); 51 | } 52 | 53 | --locked.num_open_connections; 54 | }); 55 | 56 | connectionSpotFreed(pool_key); 57 | } 58 | 59 | // Checks if the limits (global, connections open or being open by pool, or 60 | // limit per key) can fit one more connection. As a final check, checks if 61 | // it's a waste to create a new connection to avoid start opening a new 62 | // connection when we already have enough being open for the demand in 63 | // queue. If we have enough space we increment the counts. 64 | // Also, if the user supplied a permissions callback check it as well to make 65 | // sure we have permission to create a new connection 66 | bool ConnectionPoolBase::tryAddOpeningConn( 67 | const PoolKey& pool_key, 68 | std::shared_ptr context, 69 | size_t enqueued_pool_ops, 70 | uint32_t client_total_conns, 71 | uint64_t client_conn_limit, 72 | ThrottlingCallback throttlingCallback) { 73 | auto canOpen = counters_.withWLock([&](auto& locked) { 74 | if (canCreateMoreConnections( 75 | pool_key, 76 | enqueued_pool_ops, 77 | client_total_conns, 78 | client_conn_limit, 79 | locked)) { 80 | ++locked.pending_connections[pool_key]; 81 | ++locked.num_pending_connections; 82 | return true; 83 | } 84 | 85 | return false; 86 | }); 87 | 88 | if (canOpen && shouldThrottleCallback_) { 89 | canOpen = !shouldThrottleCallback_( 90 | pool_key, std::move(context), std::move(throttlingCallback)); 91 | if (!canOpen) { 92 | // TODO(jkedgar): signal that we are getting throttled so we can start 93 | // returning errors to the caller. Right now throttling will just cause 94 | // timeout errors but we want the client to know that they are throttled 95 | // ideally. 96 | 97 | // If the we didn't get permission from the callback to add more 98 | // connections, decrement the counts that got incremented above. 99 | counters_.withWLock([&](auto& locked) { 100 | --locked.pending_connections[pool_key]; 101 | --locked.num_pending_connections; 102 | }); 103 | } 104 | } 105 | 106 | return canOpen; 107 | } 108 | 109 | void ConnectionPoolBase::removeOpeningConn(const PoolKey& pool_key) { 110 | counters_.withWLock([&](auto& locked) { 111 | auto num = --locked.pending_connections[pool_key]; 112 | DCHECK_GE(int64_t(num), 0); 113 | if (num == 0) { 114 | locked.pending_connections.erase(pool_key); 115 | } 116 | --locked.num_pending_connections; 117 | }); 118 | } 119 | 120 | void ConnectionPoolBase::displayOpenConnections() { 121 | counters_.withRLock([](const auto& locked) { 122 | LOG(INFO) << "*** Open connections"; 123 | for (const auto& [key, value] : locked.open_connections) { 124 | LOG(INFO) << key.getConnectionKeyRef().getDisplayString() << ": " 125 | << value; 126 | } 127 | 128 | LOG(INFO) << "*** Pending connections"; 129 | for (const auto& [key, value] : locked.pending_connections) { 130 | LOG(INFO) << key.getConnectionKeyRef().getDisplayString() << ": " 131 | << value; 132 | } 133 | }); 134 | } 135 | 136 | int ConnectionPoolBase::getNumKeysInOpenConnections() { 137 | return counters_.rlock()->open_connections.size(); 138 | } 139 | 140 | int ConnectionPoolBase::getNumKeysInPendingConnections() { 141 | return counters_.rlock()->pending_connections.size(); 142 | } 143 | 144 | PoolKeyStats ConnectionPoolBase::getPoolKeyStats( 145 | const PoolKey& pool_key) const { 146 | return counters_.withRLock([&](const auto& locked) { 147 | PoolKeyStats stats; 148 | stats.open_connections = 149 | folly::get_default(locked.open_connections, pool_key, 0); 150 | stats.pending_connections = 151 | folly::get_default(locked.pending_connections, pool_key, 0); 152 | stats.connection_limit = perKeyLimit(); 153 | return stats; 154 | }); 155 | } 156 | 157 | } // namespace facebook::common::mysql_client 158 | -------------------------------------------------------------------------------- /squangle/mysql_client/FetchOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/FetchOperation.h" 10 | #include "squangle/mysql_client/ConnectOperation.h" 11 | #include "squangle/mysql_client/Connection.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | FetchOperation::FetchOperation( 16 | std::unique_ptr impl, 17 | std::vector&& queries) 18 | : FetchOperation(std::move(impl), MultiQuery(std::move(queries))) {} 19 | 20 | FetchOperation::FetchOperation( 21 | std::unique_ptr impl, 22 | MultiQuery&& multi_query) 23 | : queries_(std::move(multi_query)), impl_(std::move(impl)) { 24 | if (!impl_) { 25 | throw std::runtime_error("ConnectOperationImpl is null"); 26 | } 27 | 28 | impl_->setOperation(*this); 29 | } 30 | 31 | RowStream::RowStream( 32 | std::unique_ptr mysql_query_result, 33 | std::unique_ptr metadata) 34 | : mysql_query_result_(std::move(mysql_query_result)), 35 | row_fields_(std::make_shared(std::move(metadata))) {} 36 | 37 | EphemeralRow RowStream::consumeRow() { 38 | if (!current_row_.has_value()) { 39 | LOG(DFATAL) << "Illegal operation"; 40 | } 41 | EphemeralRow eph_row(std::move(*current_row_)); 42 | current_row_.reset(); 43 | return eph_row; 44 | } 45 | 46 | bool RowStream::hasNext() { 47 | // Slurp needs to happen after `consumeRow` has been called. 48 | // Because it will move the buffer. 49 | slurp(); 50 | // First iteration 51 | return current_row_.has_value(); 52 | } 53 | 54 | bool RowStream::slurp() { 55 | CHECK_THROW(mysql_query_result_ != nullptr, db::OperationStateException); 56 | if (current_row_.has_value() || query_finished_) { 57 | return true; 58 | } 59 | auto [result, row] = mysql_query_result_->fetchRow(); 60 | if (result == PENDING) { 61 | return false; 62 | } 63 | 64 | if (row == nullptr) { 65 | query_finished_ = true; 66 | return true; 67 | } 68 | current_row_.assign(EphemeralRow(std::move(row), row_fields_)); 69 | query_result_size_ += current_row_->calculateRowLength(); 70 | ++num_rows_seen_; 71 | return true; 72 | } 73 | 74 | void FetchOperationImpl::setFetchAction(FetchAction action) { 75 | if (isPaused()) { 76 | paused_action_ = action; 77 | } else { 78 | active_fetch_action_ = action; 79 | } 80 | } 81 | 82 | const MultiQuery& FetchOperationImpl::queries() const { 83 | return getOp().queries(); 84 | } 85 | 86 | const InternalConnection& FetchOperationImpl::getInternalConnection() const { 87 | return conn().getInternalConnection(); 88 | } 89 | 90 | std::string FetchOperationImpl::generateTimeoutError( 91 | std::string rowdata, 92 | Millis elapsed) const { 93 | auto cbDelay = client_.callbackDelayAvg(); 94 | bool stalled = cbDelay >= kCallbackDelayStallThreshold; 95 | 96 | std::vector parts; 97 | parts.push_back(fmt::format( 98 | "[{}]({}) Query timed out", 99 | static_cast( 100 | stalled ? SquangleErrno::SQ_ERRNO_QUERY_TIMEOUT_LOOP_STALLED 101 | : SquangleErrno::SQ_ERRNO_QUERY_TIMEOUT), 102 | kErrorPrefix)); 103 | 104 | parts.push_back(std::move(rowdata)); 105 | parts.push_back(timeoutMessage(elapsed)); 106 | if (stalled) { 107 | parts.push_back(threadOverloadMessage(cbDelay)); 108 | } 109 | 110 | return folly::join(" ", parts); 111 | } 112 | 113 | void FetchOperationImpl::cancel() { 114 | // Free any allocated results before the connection is closed 115 | // We need to do this in the mysql_thread for async versions as the 116 | // mysql_thread _might_ be using that memory 117 | auto cancelFn = [&]() { 118 | current_row_stream_ = folly::none; 119 | OperationBase::cancel(); 120 | }; 121 | if (client_.isInCorrectThread(true)) { 122 | cancelFn(); 123 | } else { 124 | client_.runInThread(std::move(cancelFn), true /*wait*/); 125 | } 126 | } 127 | 128 | uint64_t FetchOperationImpl::currentLastInsertId() const { 129 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 130 | return current_last_insert_id_; 131 | } 132 | 133 | uint64_t FetchOperationImpl::currentAffectedRows() const { 134 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 135 | return current_affected_rows_; 136 | } 137 | 138 | const std::string& FetchOperationImpl::currentRecvGtid() const { 139 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 140 | return current_recv_gtid_; 141 | } 142 | 143 | const std::optional& FetchOperationImpl::currentMysqlInfo() const { 144 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 145 | return current_mysql_info_; 146 | } 147 | 148 | const std::optional FetchOperationImpl::currentRowsMatched() const { 149 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 150 | return current_rows_matched_; 151 | } 152 | 153 | const AttributeMap& FetchOperationImpl::currentRespAttrs() const { 154 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 155 | return current_resp_attrs_; 156 | } 157 | 158 | unsigned int FetchOperationImpl::currentWarningsCount() const { 159 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 160 | return current_warnings_count_; 161 | } 162 | 163 | RowStream* FetchOperationImpl::rowStream() { 164 | CHECK_THROW(isStreamAccessAllowed(), db::OperationStateException); 165 | return current_row_stream_.get_pointer(); 166 | } 167 | 168 | AttributeMap FetchOperationImpl::readResponseAttributes() { 169 | return conn().getResponseAttributes(); 170 | } 171 | 172 | FetchOperation& FetchOperationImpl::getOp() const { 173 | DCHECK(op_ && dynamic_cast(op_) != nullptr); 174 | return *(FetchOperation*)op_; 175 | } 176 | 177 | folly::StringPiece FetchOperationImpl::toString(FetchAction action) { 178 | switch (action) { 179 | case FetchAction::StartQuery: 180 | return "StartQuery"; 181 | case FetchAction::InitFetch: 182 | return "InitFetch"; 183 | case FetchAction::Fetch: 184 | return "Fetch"; 185 | case FetchAction::WaitForConsumer: 186 | return "WaitForConsumer"; 187 | case FetchAction::CompleteQuery: 188 | return "CompleteQuery"; 189 | case FetchAction::CompleteOperation: 190 | return "CompleteOperation"; 191 | } 192 | LOG(DFATAL) << "unable to convert result to string: " 193 | << static_cast(action); 194 | return "Unknown result"; 195 | } 196 | 197 | } // namespace facebook::common::mysql_client 198 | -------------------------------------------------------------------------------- /squangle/mysql_client/Flags.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/Flags.h" 10 | 11 | DEFINE_int64( 12 | async_mysql_connect_timeout_micros, 13 | 1030 * 1000, 14 | "default timeout, in micros, for mysql connect"); 15 | 16 | DEFINE_int64( 17 | async_mysql_timeout_micros, 18 | 10 * 1000 * 1000, 19 | "default timeout, in micros, for mysql operations"); 20 | 21 | // Default timeout of 0 is no-op for connect tcp timeout 22 | DEFINE_int64( 23 | async_mysql_connect_tcp_timeout_micros, 24 | 0, 25 | "default timeout, in micros, for mysql connect"); 26 | 27 | DEFINE_int64( 28 | async_mysql_max_connect_timeout_micros, 29 | #if defined(FOLLY_SANITIZE_ADDRESS) || (FOLLY_SANITIZE_THREAD) 30 | 30 * 1000 * 1000, 31 | #else 32 | 3 * 1000 * 1000, 33 | #endif 34 | "The maximum connect timeout, to protect customers from themselves"); 35 | -------------------------------------------------------------------------------- /squangle/mysql_client/Flags.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | DECLARE_int64(async_mysql_connect_tcp_timeout_micros); 14 | DECLARE_int64(async_mysql_connect_timeout_micros); 15 | DECLARE_int64(async_mysql_max_connect_timeout_micros); 16 | DECLARE_int64(async_mysql_timeout_micros); 17 | -------------------------------------------------------------------------------- /squangle/mysql_client/InternalConnection.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/InternalConnection.h" 10 | 11 | namespace facebook::common::mysql_client { 12 | 13 | std::string InternalConnection::escapeString(std::string_view unescaped) const { 14 | std::string escaped; 15 | escaped.resize((2 * unescaped.size()) + 1); 16 | auto size = escapeString(escaped.data(), unescaped.data(), unescaped.size()); 17 | escaped.resize(size); 18 | return escaped; 19 | } 20 | 21 | } // namespace facebook::common::mysql_client 22 | -------------------------------------------------------------------------------- /squangle/mysql_client/InternalConnection.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | 15 | #include "squangle/base/Base.h" 16 | #include "squangle/mysql_client/Compression.h" 17 | #include "squangle/mysql_client/SSLOptionsProviderBase.h" 18 | 19 | namespace facebook::common::mysql_client { 20 | 21 | class ConnectionKey; 22 | class ConnectionOptions; 23 | 24 | using MysqlCertValidatorCallback = 25 | int (*)(X509* server_cert, const void* context, const char** errptr); 26 | 27 | enum InternalStatus { 28 | PENDING, 29 | DONE, 30 | ERROR, 31 | }; 32 | 33 | class InternalRow { 34 | public: 35 | enum class Type { 36 | Null, 37 | Bool, 38 | Int64, 39 | UInt64, 40 | Double, 41 | String, 42 | }; 43 | 44 | virtual ~InternalRow() = default; 45 | 46 | virtual folly::StringPiece columnString(size_t /*col*/) const { 47 | DCHECK(false) << "Not implemented"; 48 | return folly::StringPiece(); 49 | } 50 | 51 | virtual bool columnBool(size_t /*col*/) const { 52 | DCHECK(false) << "Not implemented"; 53 | return false; 54 | } 55 | 56 | virtual int64_t columnInt64(size_t /*col*/) const { 57 | DCHECK(false) << "Not implemented"; 58 | return 0; 59 | } 60 | 61 | virtual uint64_t columnUInt64(size_t /*col*/) const { 62 | DCHECK(false) << "Not implemented"; 63 | return 0; 64 | } 65 | 66 | virtual double columnDouble(size_t /*col*/) const { 67 | DCHECK(false) << "Not implemented"; 68 | return 0.0; 69 | } 70 | 71 | virtual Type columnType(size_t col) const = 0; 72 | 73 | virtual size_t columnLength(size_t col) const = 0; 74 | }; 75 | 76 | class InternalRowMetadata { 77 | public: 78 | virtual ~InternalRowMetadata() = default; 79 | 80 | virtual size_t numFields() const = 0; 81 | 82 | virtual folly::StringPiece getTableName(size_t field) const = 0; 83 | 84 | virtual folly::StringPiece getFieldName(size_t field) const = 0; 85 | 86 | virtual enum_field_types getFieldType(size_t field) const = 0; 87 | 88 | virtual uint64_t getFieldFlags(size_t field) const = 0; 89 | }; 90 | 91 | class InternalResult { 92 | public: 93 | virtual ~InternalResult() = default; 94 | 95 | using FetchRowRet = std::pair>; 96 | 97 | virtual FetchRowRet fetchRow() = 0; 98 | 99 | virtual size_t numRows() const = 0; 100 | 101 | virtual void close() = 0; 102 | 103 | // If needed this can be overridden to make sure the rows get cranked 104 | virtual folly::coro::Task co_crank() { 105 | co_return; 106 | } 107 | }; 108 | 109 | class InternalConnection { 110 | public: 111 | using Status = InternalStatus; 112 | 113 | virtual ~InternalConnection() = default; 114 | 115 | virtual void setReusable(bool /*reusable*/) {} 116 | 117 | virtual bool isReusable() const { 118 | return false; 119 | } 120 | 121 | virtual void disableCloseOnDestroy() {} 122 | 123 | virtual bool isSSL() const { 124 | return true; 125 | } 126 | 127 | virtual bool inTransaction() const = 0; 128 | 129 | virtual void setNeedResetBeforeReuse() {} 130 | 131 | virtual bool needResetBeforeReuse() { 132 | return true; 133 | } 134 | 135 | virtual std::string serverInfo() const = 0; 136 | 137 | virtual long threadId() const { 138 | return 0; 139 | } 140 | 141 | virtual void disableLocalFiles() {} 142 | 143 | virtual void disableSSL() {} 144 | 145 | virtual bool sslSessionReused() const { 146 | return false; 147 | } 148 | 149 | virtual unsigned int warningCount() const = 0; 150 | 151 | std::string escapeString(std::string_view unescaped) const; 152 | 153 | virtual size_t escapeString(char* out, const char* src, size_t length) 154 | const = 0; 155 | 156 | virtual std::function getCloseFunction() { 157 | return nullptr; 158 | } 159 | 160 | virtual unsigned int getErrno() const = 0; 161 | 162 | virtual std::string getErrorMessage() const = 0; 163 | 164 | virtual void setConnectAttributes(const AttributeMap& attributes) = 0; 165 | 166 | virtual int setQueryAttributes(const AttributeMap& attributes) = 0; 167 | 168 | virtual int setQueryAttribute( 169 | const std::string& key, 170 | const std::string& value) = 0; 171 | 172 | virtual AttributeMap getResponseAttributes() const = 0; 173 | 174 | virtual void setConnectTimeout(Millis timeout) const = 0; 175 | 176 | virtual void setReadTimeout(Millis timeout) const = 0; 177 | 178 | virtual void setWriteTimeout(Millis timeout) const = 0; 179 | 180 | virtual uint64_t getLastInsertId() const = 0; 181 | 182 | virtual uint64_t getAffectedRows() const = 0; 183 | 184 | virtual std::optional getRecvGtid() const = 0; 185 | 186 | virtual std::optional getMySQLInfo() const = 0; 187 | 188 | virtual std::optional getSchemaChanged() const = 0; 189 | 190 | virtual bool getNoIndexUsed() const = 0; 191 | 192 | virtual bool wasSlow() const = 0; 193 | 194 | virtual bool getAutocommit() const = 0; 195 | 196 | virtual bool ping() const = 0; 197 | 198 | virtual Status resetConn() const = 0; 199 | 200 | virtual Status changeUser( 201 | std::shared_ptr conn_key) const = 0; 202 | 203 | virtual bool dumpDebugInfo() const = 0; 204 | }; 205 | 206 | } // namespace facebook::common::mysql_client 207 | -------------------------------------------------------------------------------- /squangle/mysql_client/MultiQueryOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/MultiQueryOperation.h" 10 | #include "squangle/mysql_client/OperationHelpers.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | MultiQueryOperation::MultiQueryOperation( 15 | std::unique_ptr impl, 16 | std::vector&& queries) 17 | : FetchOperation(std::move(impl), std::move(queries)), 18 | current_query_result_(std::make_unique(0)) {} 19 | 20 | void MultiQueryOperation::notifyInitQuery() { 21 | auto* row_stream = rowStream(); 22 | if (row_stream) { 23 | // Populate RowFields, this is the metadata of rows. 24 | current_query_result_->setRowFields( 25 | row_stream->getEphemeralRowFields()->makeBufferedFields()); 26 | } 27 | } 28 | 29 | void MultiQueryOperation::notifyRowsReady() { 30 | // Create buffered RowBlock 31 | auto row_block = makeRowBlockFromStream( 32 | current_query_result_->getSharedRowFields(), rowStream()); 33 | if (row_block.numRows() == 0) { 34 | return; 35 | } 36 | 37 | current_query_result_->appendRowBlock(std::move(row_block)); 38 | if (buffered_query_callback_) { 39 | buffered_query_callback_( 40 | *this, current_query_result_.get(), QueryCallbackReason::RowsFetched); 41 | } 42 | } 43 | 44 | void MultiQueryOperation::notifyFailure(OperationResult result) { 45 | // This needs to be called before notifyOperationCompleted, because 46 | // in non-callback mode we "notify" the conditional variable in `Connection`. 47 | current_query_result_->setOperationResult(result); 48 | } 49 | 50 | bool MultiQueryOperation::notifyQuerySuccess(bool) { 51 | current_query_result_->setPartial(false); 52 | 53 | current_query_result_->setOperationResult(OperationResult::Succeeded); 54 | current_query_result_->setNumRowsAffected( 55 | FetchOperation::currentAffectedRows()); 56 | current_query_result_->setLastInsertId(FetchOperation::currentLastInsertId()); 57 | current_query_result_->setRecvGtid(FetchOperation::currentRecvGtid()); 58 | current_query_result_->setMysqlInfo(FetchOperation::currentMysqlInfo()); 59 | current_query_result_->setRowsMatched(FetchOperation::currentRowsMatched()); 60 | current_query_result_->setResponseAttributes( 61 | FetchOperation::currentRespAttrs()); 62 | current_query_result_->setWasSlow(FetchOperation::wasSlow()); 63 | current_query_result_->setWarningsCount( 64 | FetchOperation::currentWarningsCount()); 65 | 66 | // Notify the callback before moving the result into the operation. This is 67 | // because the callback can't access the result out of the operation. 68 | if (buffered_query_callback_) { 69 | buffered_query_callback_( 70 | *this, current_query_result_.get(), QueryCallbackReason::QueryBoundary); 71 | } 72 | 73 | query_results_.emplace_back(std::move(*current_query_result_.get())); 74 | current_query_result_ = 75 | std::make_unique(current_query_result_->queryNum() + 1); 76 | return true; 77 | } 78 | 79 | void MultiQueryOperation::notifyOperationCompleted(OperationResult result) { 80 | if (!buffered_query_callback_) { // No callback to be done 81 | return; 82 | } 83 | // Nothing that changes the non-callback state is safe to be done here. 84 | auto reason = 85 | (result == OperationResult::Succeeded ? QueryCallbackReason::Success 86 | : QueryCallbackReason::Failure); 87 | buffered_query_callback_(*this, current_query_result_.get(), reason); 88 | // Release callback since no other callbacks will be made 89 | buffered_query_callback_ = nullptr; 90 | } 91 | 92 | } // namespace facebook::common::mysql_client 93 | -------------------------------------------------------------------------------- /squangle/mysql_client/MultiQueryOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/FetchOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class MultiQueryOperation; 16 | 17 | using MultiQueryCallback = std::function< 18 | void(MultiQueryOperation&, QueryResult*, QueryCallbackReason)>; 19 | 20 | // An operation representing a query with multiple statements. 21 | // If a callback is set, it invokes the callback as rows arrive. 22 | // If there is no callback, it buffers all results into memory 23 | // and makes them available as a RowBlock. 24 | // This is inefficient for large results. 25 | // 26 | // Constructed via Connection::beginMultiQuery. 27 | class MultiQueryOperation : public FetchOperation { 28 | public: 29 | ~MultiQueryOperation() override = default; 30 | 31 | // Set our callback. This is invoked multiple times -- once for 32 | // every RowBatch and once, with nullptr for the RowBatch, 33 | // indicating the query is complete. 34 | void setCallback(MultiQueryCallback cb) { 35 | buffered_query_callback_ = std::move(cb); 36 | } 37 | void chainCallback(MultiQueryCallback cb) { 38 | auto origCb = std::move(buffered_query_callback_); 39 | if (origCb) { 40 | cb = [origCb = std::move(origCb), cb = std::move(cb)]( 41 | MultiQueryOperation& op, 42 | QueryResult* result, 43 | QueryCallbackReason reason) { 44 | origCb(op, result, reason); 45 | cb(op, result, reason); 46 | }; 47 | } 48 | setCallback(std::move(cb)); 49 | } 50 | 51 | // Steal all rows. Only valid if there is no callback. Inefficient 52 | // for large result sets. 53 | // Only call after the query has finished, don't use it inside callbacks 54 | std::vector&& stealQueryResults() { 55 | CHECK_THROW(done(), db::OperationStateException); 56 | return std::move(query_results_); 57 | } 58 | 59 | // Only call this after the query has finished and don't use it inside 60 | // callbacks 61 | const std::vector& queryResults() const { 62 | CHECK_THROW(done(), db::OperationStateException); 63 | return query_results_; 64 | } 65 | 66 | // Returns the Query for a query index. 67 | const Query& getQuery(int index) const { 68 | return queries_.getQuery(index); 69 | } 70 | 71 | // Returns the list of Queries 72 | const std::vector& getQueries() const { 73 | return queries_.getQueries(); 74 | } 75 | 76 | void setQueryResults(std::vector query_results) { 77 | query_results_ = std::move(query_results); 78 | } 79 | 80 | // Overriding to narrow the return type 81 | MultiQueryOperation& setTimeout(Duration timeout) { 82 | Operation::setTimeout(timeout); 83 | return *this; 84 | } 85 | 86 | db::OperationType getOperationType() const override { 87 | return db::OperationType::MultiQuery; 88 | } 89 | 90 | protected: 91 | void notifyInitQuery() override; 92 | void notifyRowsReady() override; 93 | bool notifyQuerySuccess(bool more_results) override; 94 | void notifyFailure(OperationResult result) override; 95 | void notifyOperationCompleted(OperationResult result) override; 96 | 97 | MultiQueryOperation( 98 | std::unique_ptr impl, 99 | std::vector&& queries); 100 | 101 | // Calls the FetchOperation specializedCompleteOperation and then does 102 | // callbacks if needed 103 | 104 | private: 105 | MultiQueryCallback buffered_query_callback_; 106 | 107 | // Storage fields for every statement in the query 108 | // Only to be used if there is no callback set. 109 | std::vector query_results_; 110 | // Buffer to trans to `query_results_` and for buffered callback. 111 | std::unique_ptr current_query_result_; 112 | 113 | int num_current_query_ = 0; 114 | 115 | friend class Connection; 116 | }; 117 | 118 | } // namespace facebook::common::mysql_client 119 | -------------------------------------------------------------------------------- /squangle/mysql_client/MultiQueryStreamOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/MultiQueryStreamOperation.h" 10 | 11 | namespace facebook::common::mysql_client { 12 | 13 | /*static*/ 14 | std::shared_ptr MultiQueryStreamOperation::create( 15 | std::unique_ptr opImpl, 16 | MultiQuery&& multi_query) { 17 | return std::shared_ptr( 18 | new MultiQueryStreamOperation(std::move(opImpl), std::move(multi_query))); 19 | } 20 | 21 | MultiQueryStreamOperation::MultiQueryStreamOperation( 22 | std::unique_ptr opImpl, 23 | MultiQuery&& multi_query) 24 | : FetchOperation(std::move(opImpl), std::move(multi_query)) {} 25 | 26 | MultiQueryStreamOperation::MultiQueryStreamOperation( 27 | std::unique_ptr opImpl, 28 | std::vector&& queries) 29 | : FetchOperation(std::move(opImpl), std::move(queries)) {} 30 | 31 | void MultiQueryStreamOperation::invokeCallback(StreamState reason) { 32 | // Construct a CallbackVistor object and pass to apply_vistor. It will 33 | // call the appropriate overaload of 'operator()' depending on the type 34 | // of callback stored in stream_callback_ i.e. either MultiQueryStreamHandler 35 | // or MultiQueryStreamOperation::Callback. 36 | std::visit(CallbackVisitor(*this, reason), stream_callback_); 37 | } 38 | 39 | void MultiQueryStreamOperation::notifyInitQuery() { 40 | invokeCallback(StreamState::InitQuery); 41 | } 42 | 43 | void MultiQueryStreamOperation::notifyRowsReady() { 44 | invokeCallback(StreamState::RowsReady); 45 | } 46 | 47 | bool MultiQueryStreamOperation::notifyQuerySuccess(bool) { 48 | // Query Boundary, only for streaming to allow the user to read from the 49 | // connection. 50 | // This will allow pause in the end of the query. End of operations don't 51 | // allow. 52 | invokeCallback(StreamState::QueryEnded); 53 | return true; 54 | } 55 | 56 | void MultiQueryStreamOperation::notifyFailure(OperationResult) { 57 | // Nop 58 | } 59 | 60 | void MultiQueryStreamOperation::notifyOperationCompleted( 61 | OperationResult result) { 62 | auto reason = 63 | (result == OperationResult::Succeeded ? StreamState::Success 64 | : StreamState::Failure); 65 | 66 | invokeCallback(reason); 67 | stream_callback_ = nullptr; 68 | } 69 | 70 | } // namespace facebook::common::mysql_client 71 | -------------------------------------------------------------------------------- /squangle/mysql_client/MultiQueryStreamOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/FetchOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class MultiQueryStreamHandler; 16 | class SyncConnection; 17 | 18 | // This operation only supports one mode: streaming callback. This is a 19 | // simple layer on top of FetchOperation to adapt from `notify` to 20 | // StreamCallback. 21 | // This is an experimental class. Please don't use directly. 22 | class MultiQueryStreamOperation : public FetchOperation { 23 | public: 24 | ~MultiQueryStreamOperation() override = default; 25 | 26 | using Callback = std::function; 27 | using StreamCallback = std::variant; 28 | 29 | void notifyInitQuery() override; 30 | void notifyRowsReady() override; 31 | bool notifyQuerySuccess(bool more_results) override; 32 | void notifyFailure(OperationResult result) override; 33 | void notifyOperationCompleted(OperationResult result) override; 34 | 35 | // Overriding to narrow the return type 36 | MultiQueryStreamOperation& setTimeout(Duration timeout) { 37 | Operation::setTimeout(timeout); 38 | return *this; 39 | } 40 | 41 | db::OperationType getOperationType() const override { 42 | return db::OperationType::MultiQueryStream; 43 | } 44 | 45 | void setCallback(StreamCallback cb) { 46 | stream_callback_ = std::move(cb); 47 | } 48 | 49 | protected: 50 | static std::shared_ptr create( 51 | std::unique_ptr opImpl, 52 | MultiQuery&& multi_query); 53 | 54 | friend Connection; 55 | friend SyncConnection; 56 | 57 | private: 58 | MultiQueryStreamOperation( 59 | std::unique_ptr opImpl, 60 | MultiQuery&& multi_query); 61 | MultiQueryStreamOperation( 62 | std::unique_ptr opImpl, 63 | std::vector&& queries); 64 | 65 | // wrapper to construct CallbackVistor and invoke the 66 | // right callback 67 | void invokeCallback(StreamState state); 68 | 69 | // Vistor to invoke the right callback depending on the type stored 70 | // in the variant 'stream_callback_' 71 | struct CallbackVisitor { 72 | CallbackVisitor(MultiQueryStreamOperation& op, StreamState state) 73 | : op_(op), state_(state) {} 74 | 75 | void operator()(MultiQueryStreamHandler* handler) const { 76 | if (handler != nullptr) { 77 | handler->streamCallback(op_, state_); 78 | } 79 | } 80 | 81 | void operator()(const Callback& cb) const { 82 | if (cb != nullptr) { 83 | cb(op_, state_); 84 | } 85 | } 86 | 87 | private: 88 | MultiQueryStreamOperation& op_; 89 | StreamState state_; 90 | }; 91 | 92 | StreamCallback stream_callback_; 93 | }; 94 | 95 | } // namespace facebook::common::mysql_client 96 | -------------------------------------------------------------------------------- /squangle/mysql_client/MysqlClientBase.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/MysqlClientBase.h" 10 | #include "squangle/mysql_client/ConnectOperation.h" 11 | #include "squangle/mysql_client/Connection.h" 12 | #include "squangle/mysql_client/SpecialOperation.h" 13 | 14 | namespace { 15 | 16 | // Used to initialize the SSL and MySQL libraries once per binary 17 | class InitMysqlLibrary { 18 | public: 19 | InitMysqlLibrary() { 20 | mysql_library_init(-1, nullptr, nullptr); 21 | } 22 | ~InitMysqlLibrary() { 23 | mysql_library_end(); 24 | } 25 | 26 | // don't allow copying or moving 27 | InitMysqlLibrary(const InitMysqlLibrary&) = delete; 28 | InitMysqlLibrary& operator=(const InitMysqlLibrary&) = delete; 29 | 30 | InitMysqlLibrary(InitMysqlLibrary&&) = delete; 31 | InitMysqlLibrary& operator=(InitMysqlLibrary&&) = delete; 32 | }; 33 | 34 | } // namespace 35 | 36 | namespace facebook::common::mysql_client { 37 | 38 | // mysql_library_init() and mysql_library_end() need to run on the same thread 39 | [[maybe_unused]] static InitMysqlLibrary unused; 40 | 41 | MysqlClientBase::MysqlClientBase( 42 | std::unique_ptr db_logger, 43 | std::unique_ptr db_stats, 44 | std::unique_ptr exception_builder) 45 | : db_logger_(std::move(db_logger)), 46 | client_stats_(std::move(db_stats)), 47 | exception_builder_(std::move(exception_builder)) { 48 | if (!exception_builder_) { 49 | exception_builder_ = std::make_unique(); 50 | } 51 | } 52 | 53 | void MysqlClientBase::logQuerySuccess( 54 | const db::QueryLoggingData& logging_data, 55 | const Connection& conn) { 56 | auto conn_context = conn.getConnectionContext(); 57 | stats()->incrSucceededQueries(conn_context); 58 | 59 | if (db_logger_) { 60 | db_logger_->logQuerySuccess( 61 | logging_data, makeSquangleLoggingData(conn.getKey(), conn_context)); 62 | } 63 | } 64 | 65 | void MysqlClientBase::logQueryFailure( 66 | const db::QueryLoggingData& logging_data, 67 | db::FailureReason reason, 68 | unsigned int mysqlErrno, 69 | const std::string& error, 70 | const Connection& conn) { 71 | auto conn_context = conn.getConnectionContext(); 72 | stats()->incrFailedQueries(conn_context, mysqlErrno, error); 73 | 74 | if (db_logger_) { 75 | db_logger_->logQueryFailure( 76 | logging_data, 77 | reason, 78 | mysqlErrno, 79 | error, 80 | makeSquangleLoggingData(conn.getKey(), conn_context)); 81 | } 82 | } 83 | 84 | void MysqlClientBase::logConnectionSuccess( 85 | const db::CommonLoggingData& logging_data, 86 | std::shared_ptr conn_key, 87 | const db::ConnectionContextBase* connection_context) { 88 | if (db_logger_) { 89 | db_logger_->logConnectionSuccess( 90 | logging_data, 91 | makeSquangleLoggingData(std::move(conn_key), connection_context)); 92 | } 93 | } 94 | 95 | void MysqlClientBase::logConnectionFailure( 96 | const db::CommonLoggingData& logging_data, 97 | db::FailureReason reason, 98 | std::shared_ptr conn_key, 99 | unsigned int mysqlErrno, 100 | const std::string& error, 101 | const db::ConnectionContextBase* connection_context) { 102 | stats()->incrFailedConnections(connection_context, mysqlErrno, error); 103 | 104 | if (db_logger_) { 105 | db_logger_->logConnectionFailure( 106 | logging_data, 107 | reason, 108 | mysqlErrno, 109 | error, 110 | makeSquangleLoggingData(std::move(conn_key), connection_context)); 111 | } 112 | } 113 | 114 | std::shared_ptr MysqlClientBase::beginConnection( 115 | const std::string& host, 116 | int port, 117 | const std::string& database_name, 118 | const std::string& user, 119 | const std::string& password) { 120 | return beginConnection(std::make_shared( 121 | host, port, database_name, user, password)); 122 | } 123 | 124 | std::shared_ptr MysqlClientBase::beginConnection( 125 | std::shared_ptr conn_key) { 126 | auto impl = createConnectOperationImpl(this, std::move(conn_key)); 127 | auto ret = ConnectOperation::create(std::move(impl)); 128 | if (connection_cb_) { 129 | ret->setObserverCallback(connection_cb_); 130 | } 131 | return ret; 132 | } 133 | 134 | // Helper versions of the above that take a Connection instead of a 135 | // ConnectionProxy 136 | std::unique_ptr MysqlClientBase::createFetchOperationImpl( 137 | std::unique_ptr conn) const { 138 | return createFetchOperationImpl( 139 | std::make_unique(std::move(conn))); 140 | } 141 | std::unique_ptr 142 | MysqlClientBase::createSpecialOperationImpl( 143 | std::unique_ptr conn) const { 144 | return createSpecialOperationImpl( 145 | std::make_unique(std::move(conn))); 146 | } 147 | 148 | } // namespace facebook::common::mysql_client 149 | -------------------------------------------------------------------------------- /squangle/mysql_client/MysqlClientBase.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | #include "squangle/base/ConnectionKey.h" 14 | #include "squangle/logger/DBEventLogger.h" 15 | #include "squangle/mysql_client/MysqlExceptionBuilder.h" 16 | #include "squangle/mysql_client/Operation.h" 17 | 18 | namespace facebook::common::mysql_client { 19 | 20 | class Connection; 21 | class ConnectOperation; 22 | class ConnectOperationImpl; 23 | template 24 | class ConnectPoolOperationImpl; 25 | class FetchOperationImpl; 26 | class SpecialOperationImpl; 27 | 28 | class MysqlClientBase { 29 | public: 30 | virtual ~MysqlClientBase() = default; 31 | 32 | // Initiate a connection to a database. This is the main entrypoint. 33 | virtual std::shared_ptr beginConnection( 34 | const std::string& host, 35 | int port, 36 | const std::string& database_name, 37 | const std::string& user, 38 | const std::string& password); 39 | 40 | std::shared_ptr beginConnection( 41 | std::shared_ptr conn_key); 42 | 43 | // Factory method 44 | virtual std::unique_ptr createConnection( 45 | std::shared_ptr conn_key) = 0; 46 | 47 | virtual folly::EventBase* getEventBase() { 48 | return nullptr; 49 | } 50 | 51 | virtual const folly::EventBase* getEventBase() const { 52 | return nullptr; 53 | } 54 | 55 | void logQuerySuccess( 56 | const db::QueryLoggingData& logging_data, 57 | const Connection& conn); 58 | 59 | void logQueryFailure( 60 | const db::QueryLoggingData& logging_data, 61 | db::FailureReason reason, 62 | unsigned int mysqlErrno, 63 | const std::string& error, 64 | const Connection& conn); 65 | 66 | void logConnectionSuccess( 67 | const db::CommonLoggingData& logging_data, 68 | std::shared_ptr conn_key, 69 | const db::ConnectionContextBase* extra_logging_data); 70 | 71 | void logConnectionFailure( 72 | const db::CommonLoggingData& logging_data, 73 | db::FailureReason reason, 74 | std::shared_ptr conn_key, 75 | unsigned int mysqlErrno, 76 | const std::string& error, 77 | const db::ConnectionContextBase* extra_logging_data); 78 | 79 | db::DBCounterBase* stats() { 80 | return client_stats_.get(); 81 | } 82 | db::SquangleLoggerBase* dbLogger() { 83 | return db_logger_.get(); 84 | } 85 | const MysqlExceptionBuilder& exceptionBuilder() { 86 | return *exception_builder_; 87 | } 88 | 89 | // For internal (testing) use only 90 | std::unique_ptr setDBLoggerForTesting( 91 | std::unique_ptr dbLogger) { 92 | std::swap(db_logger_, dbLogger); 93 | return dbLogger; 94 | } 95 | std::unique_ptr setDBCounterForTesting( 96 | std::unique_ptr dbCounter) { 97 | std::swap(client_stats_, dbCounter); 98 | return dbCounter; 99 | } 100 | 101 | void setConnectionCallback(ObserverCallback connection_cb) { 102 | if (connection_cb_) { 103 | connection_cb_ = [old_cb = connection_cb_, 104 | new_cb = std::move(connection_cb)](Operation& op) { 105 | old_cb(op); 106 | new_cb(op); 107 | }; 108 | } else { 109 | connection_cb_ = std::move(connection_cb); 110 | } 111 | } 112 | 113 | explicit MysqlClientBase( 114 | std::unique_ptr db_logger = nullptr, 115 | std::unique_ptr db_stats = 116 | std::make_unique(), 117 | std::unique_ptr exception_builder = nullptr); 118 | 119 | virtual bool runInThread(std::function&& fn, bool /*wait*/ = false) { 120 | fn(); 121 | return true; 122 | } 123 | 124 | virtual uint32_t numStartedAndOpenConnections() { 125 | return 0; 126 | } 127 | virtual Duration callbackDelayAvg() const { 128 | return Duration(0); 129 | } 130 | 131 | virtual bool supportsLocalFiles() = 0; 132 | 133 | static constexpr bool implementsPooling() { 134 | return false; 135 | } 136 | 137 | virtual std::unique_ptr createConnectOperationImpl( 138 | MysqlClientBase* client, 139 | std::shared_ptr conn_key) const = 0; 140 | virtual std::unique_ptr createFetchOperationImpl( 141 | std::unique_ptr conn) const = 0; 142 | virtual std::unique_ptr createSpecialOperationImpl( 143 | std::unique_ptr conn) const = 0; 144 | 145 | // Helper versions of the above that take a Connection instead of a 146 | // ConnectionProxy 147 | std::unique_ptr createFetchOperationImpl( 148 | std::unique_ptr conn) const; 149 | std::unique_ptr createSpecialOperationImpl( 150 | std::unique_ptr conn) const; 151 | 152 | virtual void activeConnectionAdded( 153 | std::shared_ptr /*key*/) {} 154 | virtual void activeConnectionRemoved( 155 | std::shared_ptr /*key*/) {} 156 | 157 | virtual bool isInCorrectThread(bool /*expectMysqlThread*/) const { 158 | return true; 159 | } 160 | 161 | protected: 162 | friend class Connection; 163 | friend class OperationBase; 164 | friend class OperationImpl; 165 | friend class ConnectOperationImpl; 166 | template 167 | friend class ConnectPoolOperation; 168 | template 169 | friend class ConnectionPool; 170 | friend class FetchOperationImpl; 171 | friend class SpecialOperationImpl; 172 | friend class ResetOperation; 173 | friend class ChangeUserOperation; 174 | friend class ConnectionHolder; 175 | friend class AsyncConnection; 176 | friend class SyncConnection; 177 | virtual db::SquangleLoggingData makeSquangleLoggingData( 178 | std::shared_ptr connKey, 179 | const db::ConnectionContextBase* connContext) { 180 | return db::SquangleLoggingData(std::move(connKey), connContext); 181 | } 182 | 183 | virtual void addOperation(std::shared_ptr /*op*/) {} 184 | virtual void deferRemoveOperation(Operation* /*op*/) {} 185 | 186 | // Using unique pointer due inheritance virtual calls 187 | std::unique_ptr db_logger_; 188 | std::unique_ptr client_stats_; 189 | ObserverCallback connection_cb_; 190 | std::unique_ptr exception_builder_; 191 | }; 192 | 193 | } // namespace facebook::common::mysql_client 194 | -------------------------------------------------------------------------------- /squangle/mysql_client/MysqlExceptionBuilder.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | #include "squangle/base/ConnectionKey.h" 11 | #include "squangle/mysql_client/Operation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class MysqlExceptionBuilder { 16 | public: 17 | virtual ~MysqlExceptionBuilder() = default; 18 | 19 | virtual folly::exception_wrapper buildMysqlException( 20 | OperationResult result, 21 | unsigned int mysql_errno, 22 | const std::string& mysql_error, 23 | const std::shared_ptr& conn_key, 24 | Duration elapsed_time) const { 25 | return folly::make_exception_wrapper( 26 | result, mysql_errno, mysql_error, conn_key, elapsed_time); 27 | } 28 | 29 | virtual folly::exception_wrapper buildQueryException( 30 | int num_executed_queries, 31 | OperationResult result, 32 | unsigned int mysql_errno, 33 | const std::string& mysql_error, 34 | const std::shared_ptr& conn_key, 35 | Duration elapsed_time) const { 36 | return folly::make_exception_wrapper( 37 | num_executed_queries, 38 | result, 39 | mysql_errno, 40 | mysql_error, 41 | conn_key, 42 | elapsed_time); 43 | } 44 | }; 45 | } // namespace facebook::common::mysql_client 46 | -------------------------------------------------------------------------------- /squangle/mysql_client/OperationBatch.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | // 10 | #include "squangle/mysql_client/OperationBatch.h" 11 | 12 | #include 13 | 14 | namespace facebook::common::mysql_client { 15 | 16 | void OperationBatch::add(std::shared_ptr op) { 17 | CHECK_THROW( 18 | op->state() == OperationState::Unstarted, db::OperationStateException); 19 | { 20 | std::lock_guard lock(*mutex_); 21 | num_living_operations_++; 22 | } 23 | 24 | op->setObserverCallback([this](Operation& op) { 25 | std::lock_guard lock(*mutex_); 26 | if (!op.ok()) { 27 | successful_ = false; 28 | } 29 | if (--num_living_operations_ == 0) { 30 | currently_idle_->notify_one(); 31 | } 32 | }); 33 | } 34 | 35 | /* Only the creator of the OperationBatch can call drain on it */ 36 | void OperationBatch::drain() { 37 | // Drain (and hence our destructor) should be a no-op if mutex_ was 38 | // std::move'd away (which indicates this batch is the withered husk 39 | // of a std::move'd batch). 40 | if (!mutex_) { 41 | return; 42 | } 43 | 44 | DCHECK_EQ(std::this_thread::get_id(), creator_thread_id_); 45 | 46 | std::unique_lock lock(*mutex_); 47 | 48 | // no operation to wait on 49 | if (num_living_operations_ == 0) { 50 | return; 51 | } 52 | 53 | // Now wait for the operations to complete. 54 | currently_idle_->wait(lock, [this] { return (num_living_operations_ == 0); }); 55 | } 56 | 57 | bool OperationBatch::ok() { 58 | std::lock_guard lock(*mutex_); 59 | return num_living_operations_ == 0 && successful_; 60 | } 61 | 62 | void OperationBatch::markFailure() { 63 | std::lock_guard lock(*mutex_); 64 | successful_ = false; 65 | } 66 | 67 | } // namespace facebook::common::mysql_client 68 | -------------------------------------------------------------------------------- /squangle/mysql_client/OperationBatch.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | // 10 | // The OperationBatch is a class that allows waiting only for a 11 | // certain set of operations. All the operations (Connect, Query) 12 | // must be added to the batch and then drain() can be called to 13 | // wait for completion. 14 | // 15 | // It is useful for parallel queries. Because usually 16 | // the FbAsyncMysqlClient::defaultClient() will be used and it is 17 | // shared among all the threads, calling drain on the client is not 18 | // an option. This class allows waiting only on the operations started 19 | // by that thread. 20 | // An usage example can be found in the TestOperationBatch UNITTEST 21 | // from the AsyncMysqlTest.cpp file. 22 | 23 | #pragma once 24 | 25 | #include 26 | #include 27 | 28 | // #include "squangle/mysql_client/AsyncMysqlClient.h" 29 | #include "squangle/mysql_client/Operation.h" 30 | 31 | namespace facebook::common::mysql_client { 32 | 33 | class OperationBatch { 34 | public: 35 | explicit OperationBatch() 36 | : mutex_(new std::mutex), 37 | currently_idle_(new std::condition_variable), 38 | num_living_operations_(0), 39 | creator_thread_id_(std::this_thread::get_id()), 40 | successful_(true) {} 41 | 42 | // not copyable 43 | OperationBatch(const OperationBatch& other) = delete; 44 | OperationBatch& operator=(const OperationBatch& other) = delete; 45 | 46 | // movable 47 | OperationBatch(OperationBatch&& other) = default; 48 | OperationBatch& operator=(OperationBatch&& other) = default; 49 | 50 | ~OperationBatch() { 51 | drain(); 52 | } 53 | 54 | void add(std::shared_ptr op); 55 | 56 | void drain(); 57 | 58 | // Returns false if there was a failure - either an operation failed or 59 | // it was marked manually with markFailure() 60 | bool ok(); 61 | 62 | // Offers the possibility to mark from a callback that the batch was not 63 | // successful 64 | // If any operation failed, it will be marked automatically as failed. 65 | void markFailure(); 66 | 67 | private: 68 | // mutex_ is used by currently_idle_ condition variable and to protect 69 | // num_living_operations_ and successful_ variables 70 | std::unique_ptr mutex_; 71 | std::unique_ptr currently_idle_; 72 | 73 | // Counter for the number of living operations 74 | // This is used for draining ; 75 | uint32_t num_living_operations_; 76 | 77 | std::thread::id creator_thread_id_; 78 | 79 | // Indicator of the success of the batch 80 | bool successful_; 81 | }; 82 | 83 | } // namespace facebook::common::mysql_client 84 | -------------------------------------------------------------------------------- /squangle/mysql_client/OperationHelpers.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/OperationHelpers.h" 10 | 11 | namespace facebook::common::mysql_client { 12 | 13 | namespace { 14 | 15 | void copyRowToRowBlock(RowBlock* block, const EphemeralRow& eph_row) { 16 | block->startRow(); 17 | for (int i = 0; i < eph_row.numFields(); ++i) { 18 | switch (eph_row.getType(i)) { 19 | case InternalRow::Type::Null: 20 | block->appendNull(); 21 | break; 22 | case InternalRow::Type::Bool: 23 | block->appendValue(eph_row.getBool(i)); 24 | break; 25 | 26 | case InternalRow::Type::Int64: 27 | block->appendValue(eph_row.getInt64(i)); 28 | break; 29 | 30 | case InternalRow::Type::UInt64: 31 | block->appendValue(eph_row.getUInt64(i)); 32 | break; 33 | 34 | case InternalRow::Type::Double: 35 | block->appendValue(eph_row.getDouble(i)); 36 | break; 37 | 38 | case InternalRow::Type::String: 39 | block->appendValue(eph_row.getString(i)); 40 | break; 41 | } 42 | } 43 | block->finishRow(); 44 | } 45 | 46 | } // namespace 47 | 48 | RowBlock makeRowBlockFromStream( 49 | std::shared_ptr row_fields, 50 | RowStream* row_stream) { 51 | RowBlock row_block(std::move(row_fields)); 52 | // Consume row_stream 53 | while (row_stream->hasNext()) { 54 | auto eph_row = row_stream->consumeRow(); 55 | copyRowToRowBlock(&row_block, eph_row); 56 | } 57 | return row_block; 58 | } 59 | 60 | } // namespace facebook::common::mysql_client 61 | -------------------------------------------------------------------------------- /squangle/mysql_client/OperationHelpers.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/FetchOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | RowBlock makeRowBlockFromStream( 16 | std::shared_ptr row_fields, 17 | RowStream* row_stream); 18 | 19 | } // namespace facebook::common::mysql_client 20 | -------------------------------------------------------------------------------- /squangle/mysql_client/PoolKey.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/base/ConnectionKey.h" 12 | #include "squangle/mysql_client/ConnectionOptions.h" 13 | 14 | namespace facebook::common::mysql_client { 15 | 16 | class PoolKey { 17 | public: 18 | // Hashes Connections and Operations waiting for connections based on basic 19 | // Connection info (ConnectionKey) and Connection Attributes. 20 | PoolKey( 21 | std::shared_ptr conn_key, 22 | ConnectionOptions conn_opts) 23 | : connKey_(std::move(conn_key)), connOptions_(std::move(conn_opts)) { 24 | options_hash_ = folly::hash::hash_range( 25 | connOptions_.getAttributes().begin(), 26 | connOptions_.getAttributes().end()); 27 | partial_hash_ = 28 | folly::hash::hash_combine(connKey_->partial_hash(), options_hash_); 29 | full_hash_ = folly::hash::hash_combine(connKey_->hash(), options_hash_); 30 | } 31 | 32 | [[nodiscard]] bool operator==(const PoolKey& rhs) const noexcept { 33 | return full_hash_ == rhs.full_hash_ && options_hash_ == rhs.options_hash_ && 34 | *connKey_ == *rhs.connKey_; 35 | } 36 | 37 | [[nodiscard]] bool operator!=(const PoolKey& rhs) const noexcept { 38 | return !(*this == rhs); 39 | } 40 | 41 | [[nodiscard]] bool partialCompare(const PoolKey& rhs) const noexcept { 42 | return partial_hash_ == rhs.partial_hash_ && 43 | options_hash_ == rhs.options_hash_ && 44 | connKey_->partialEqual(*rhs.connKey_); 45 | } 46 | 47 | [[nodiscard]] std::shared_ptr getConnectionKey() 48 | const noexcept { 49 | return connKey_; 50 | } 51 | 52 | [[nodiscard]] const ConnectionKey& getConnectionKeyRef() const noexcept { 53 | return *connKey_; 54 | } 55 | 56 | [[nodiscard]] const ConnectionOptions& getConnectionOptions() const noexcept { 57 | return connOptions_; 58 | } 59 | 60 | [[nodiscard]] size_t getHash() const noexcept { 61 | return full_hash_; 62 | } 63 | 64 | [[nodiscard]] size_t getPartialHash() const noexcept { 65 | return partial_hash_; 66 | } 67 | 68 | [[nodiscard]] size_t getOptionsHash() const noexcept { 69 | return options_hash_; 70 | } 71 | 72 | private: 73 | std::shared_ptr connKey_; 74 | ConnectionOptions connOptions_; 75 | 76 | size_t options_hash_; 77 | size_t full_hash_; 78 | size_t partial_hash_; 79 | }; 80 | 81 | struct PoolKeyStats { 82 | size_t open_connections; 83 | size_t pending_connections; 84 | size_t connection_limit; 85 | }; 86 | 87 | std::ostream& operator<<(std::ostream& os, const PoolKey& key); 88 | 89 | class PoolKeyHash { 90 | public: 91 | size_t operator()(const PoolKey& k) const { 92 | return k.getHash(); 93 | } 94 | }; 95 | 96 | class PoolKeyPartialHash { 97 | public: 98 | size_t operator()(const PoolKey& k) const { 99 | return k.getPartialHash(); 100 | } 101 | 102 | bool operator()(const PoolKey& lhs, const PoolKey& rhs) const { 103 | return lhs.partialCompare(rhs); 104 | } 105 | }; 106 | 107 | } // namespace facebook::common::mysql_client 108 | -------------------------------------------------------------------------------- /squangle/mysql_client/QueryGenerator.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/Query.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class QueryGenerator { 16 | public: 17 | virtual ~QueryGenerator() = default; 18 | virtual Query query() = 0; 19 | }; 20 | 21 | } // namespace facebook::common::mysql_client 22 | -------------------------------------------------------------------------------- /squangle/mysql_client/QueryOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/QueryOperation.h" 10 | #include "squangle/mysql_client/OperationHelpers.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | QueryOperation::QueryOperation( 15 | std::unique_ptr impl, 16 | Query&& query) 17 | : FetchOperation(std::move(impl), std::vector{std::move(query)}), 18 | query_result_(std::make_unique(0)) {} 19 | 20 | void QueryOperation::notifyInitQuery() { 21 | auto* row_stream = rowStream(); 22 | if (row_stream) { 23 | // Populate RowFields, this is the metadata of rows. 24 | query_result_->setRowFields( 25 | row_stream->getEphemeralRowFields()->makeBufferedFields()); 26 | } 27 | } 28 | 29 | void QueryOperation::notifyRowsReady() { 30 | // QueryOperation acts as consumer of FetchOperation, and will buffer the 31 | // result. 32 | auto row_block = 33 | makeRowBlockFromStream(query_result_->getSharedRowFields(), rowStream()); 34 | 35 | // Empty result set 36 | if (row_block.numRows() == 0) { 37 | return; 38 | } 39 | 40 | query_result_->appendRowBlock(std::move(row_block)); 41 | if (buffered_query_callback_) { 42 | buffered_query_callback_( 43 | *this, query_result_.get(), QueryCallbackReason::RowsFetched); 44 | } 45 | } 46 | 47 | bool QueryOperation::notifyQuerySuccess(bool more_results) { 48 | if (more_results) { 49 | // This is the single-query API; we can't support multi-queries here, so if 50 | // we have more results we need to generate an error and cancel. 51 | setAsyncClientError( 52 | (unsigned int)SquangleErrno::SQ_INVALID_API_USAGE, 53 | "Multi-queries are not supported in this API - " 54 | "use the multi-query API instead"); 55 | cancel(); 56 | return false; 57 | } 58 | 59 | query_result_->setOperationResult(OperationResult::Succeeded); 60 | query_result_->setNumRowsAffected(FetchOperation::currentAffectedRows()); 61 | query_result_->setLastInsertId(FetchOperation::currentLastInsertId()); 62 | query_result_->setRecvGtid(FetchOperation::currentRecvGtid()); 63 | query_result_->setMysqlInfo(FetchOperation::currentMysqlInfo()); 64 | query_result_->setRowsMatched(FetchOperation::currentRowsMatched()); 65 | query_result_->setWasSlow(FetchOperation::wasSlow()); 66 | query_result_->setResponseAttributes(FetchOperation::currentRespAttrs()); 67 | query_result_->setWarningsCount(FetchOperation::currentWarningsCount()); 68 | 69 | query_result_->setPartial(false); 70 | 71 | // We are not going to make callback to user now since this only one query, 72 | // we make when we finish the operation 73 | 74 | return true; 75 | } 76 | 77 | void QueryOperation::notifyFailure(OperationResult result) { 78 | // Next call will be to notify user 79 | query_result_->setOperationResult(result); 80 | } 81 | 82 | void QueryOperation::notifyOperationCompleted(OperationResult result) { 83 | if (!buffered_query_callback_) { 84 | return; 85 | } 86 | 87 | // Nothing that changes the non-callback state is safe to be done here. 88 | auto reason = 89 | (result == OperationResult::Succeeded ? QueryCallbackReason::Success 90 | : QueryCallbackReason::Failure); 91 | buffered_query_callback_(*this, query_result_.get(), reason); 92 | // Release callback since no other callbacks will be made 93 | buffered_query_callback_ = nullptr; 94 | } 95 | 96 | } // namespace facebook::common::mysql_client 97 | -------------------------------------------------------------------------------- /squangle/mysql_client/QueryOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/FetchOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | using QueryCallback = 16 | std::function; 17 | 18 | // An operation representing a query. If a callback is set, it 19 | // invokes the callback as rows arrive. If there is no callback, it 20 | // buffers all results into memory and makes them available as a 21 | // RowBlock. This is inefficient for large results. 22 | // 23 | // Constructed via Connection::beginQuery. 24 | class QueryOperation : public FetchOperation { 25 | public: 26 | ~QueryOperation() override = default; 27 | 28 | void setCallback(QueryCallback cb) { 29 | buffered_query_callback_ = std::move(cb); 30 | } 31 | void chainCallback(QueryCallback cb) { 32 | if (buffered_query_callback_) { 33 | cb = [old_cb = buffered_query_callback_, new_cb = std::move(cb)]( 34 | QueryOperation& op, 35 | QueryResult* result, 36 | QueryCallbackReason reason) { 37 | old_cb(op, result, reason); 38 | new_cb(op, result, reason); 39 | }; 40 | } 41 | setCallback(std::move(cb)); 42 | } 43 | 44 | // Steal all rows. Only valid if there is no callback. Inefficient 45 | // for large result sets. 46 | QueryResult&& stealQueryResult() { 47 | CHECK_THROW(ok(), db::OperationStateException); 48 | return std::move(*query_result_); 49 | } 50 | 51 | const QueryResult& queryResult() const { 52 | CHECK_THROW(ok(), db::OperationStateException); 53 | return *query_result_; 54 | } 55 | 56 | // Returns the Query of this operation 57 | const Query& getQuery() const { 58 | return queries_.getQuery(0); 59 | } 60 | 61 | // Steal all rows. Only valid if there is no callback. Inefficient 62 | // for large result sets. 63 | std::vector&& stealRows() { 64 | return query_result_->stealRows(); 65 | } 66 | 67 | const std::vector& rows() const { 68 | return query_result_->rows(); 69 | } 70 | 71 | // Last insert id (aka mysql_insert_id). 72 | uint64_t lastInsertId() const { 73 | return query_result_->lastInsertId(); 74 | } 75 | 76 | // Number of rows affected (aka mysql_affected_rows). 77 | uint64_t numRowsAffected() const { 78 | return query_result_->numRowsAffected(); 79 | } 80 | 81 | // Received gtid. 82 | const std::string& recvGtid() const { 83 | return query_result_->recvGtid(); 84 | } 85 | 86 | const std::optional& mysqlInfo() const { 87 | return query_result_->mysqlInfo(); 88 | } 89 | 90 | const std::optional& rowsMatched() const { 91 | return query_result_->rowsMatched(); 92 | } 93 | 94 | void setQueryResult(QueryResult query_result) { 95 | query_result_ = std::make_unique(std::move(query_result)); 96 | } 97 | 98 | // Overriding to narrow the return type 99 | QueryOperation& setTimeout(Duration timeout) { 100 | Operation::setTimeout(timeout); 101 | return *this; 102 | } 103 | 104 | db::OperationType getOperationType() const override { 105 | return db::OperationType::Query; 106 | } 107 | 108 | protected: 109 | void notifyInitQuery() override; 110 | void notifyRowsReady() override; 111 | bool notifyQuerySuccess(bool more_results) override; 112 | void notifyFailure(OperationResult result) override; 113 | void notifyOperationCompleted(OperationResult result) override; 114 | 115 | QueryOperation(std::unique_ptr impl, Query&& query); 116 | 117 | private: 118 | QueryCallback buffered_query_callback_; 119 | std::unique_ptr query_result_; 120 | friend class Connection; 121 | }; 122 | 123 | } // namespace facebook::common::mysql_client 124 | -------------------------------------------------------------------------------- /squangle/mysql_client/README.md: -------------------------------------------------------------------------------- 1 | squangle/mysql_client 2 | -------------------- 3 | 4 | facebook::common::mysql_client is a collection of classes for 5 | performing fully asynchronous queries against MySQL databases. It 6 | supports connecting and querying in asynchronous modes (either 7 | callbacks or explicit waiting). 8 | 9 | This document is a basic overview/roadmap; please consult the headers 10 | for more details, as well as the test cases and ParallelMysqlTool.cpp 11 | for example usage. 12 | 13 | #### Synchronous Examples 14 | ### Basic Example 15 | 16 | #include "squangle/mysql_client/AsyncMysqlClient.h" 17 | 18 | AsyncMysqlClient* client = AsyncMysqlClient::defaultClient(); 19 | 20 | std::shared_ptr connect_op = 21 | client->beginConnection(hostname, port, dbname, user, password); 22 | connect_op->run(); 23 | connect_op->wait(); 24 | if (!connect_op.ok()) { 25 | panic(); 26 | } 27 | 28 | ### Connect and Query Example 29 | 30 | Query query("SELECT %C FROM %T WHERE %C = %s AND %C BETWEEN %d AND %d", 31 | "pusheen_macro", 32 | "macros", 33 | "type", "pusheen", 34 | "awesomeness", MAX, INF); 35 | 36 | std::shared_ptr connect_op = 37 | client->beginConnection(hostname, port, dbname, user, password); 38 | conn_op->run()->wait(); 39 | if (!conn_op->ok()) { 40 | panicalittlebit(); 41 | } 42 | 43 | auto query_op = Connection::beginQuery(conn_op->releaseConnection(), std::move(query)); 44 | query_op->run()->wait(); 45 | if (!query_op->ok()) { 46 | dontpanicthatmuch(); 47 | } 48 | 49 | for (const auto& row : query_op->queryResult()) { 50 | // Play with rows :) 51 | } 52 | 53 | #### Asynchronous Examples 54 | ### Single Query Example 55 | 56 | std::shared_ptr query_op = 57 | Connection::beginQuery(std::move(conn), query_string); 58 | query_op->setCallback(query_callback); 59 | query_op->setTimeout(500ms); 60 | query_op->run(); 61 | 62 | void query_callback(QueryOperation& op, 63 | QueryResult* query_result, 64 | QueryCallbackReason reason) { 65 | if (reason == QueryCallbackReason::RowsFetched) { 66 | LOG(INFO) << "Saw " << query_result->numRows() << " more rows!"; 67 | for (const auto& row : *query_result) { 68 | LOG(INFO) << "Row: " << folly::join(row, "\t"); 69 | } 70 | } else if (reason == QueryCallbackReason::Success) { 71 | LOG(INFO) << "Query succeeded!"; 72 | } else { 73 | LOG(ERROR) << "Query failed! " << op.mysql_error(); 74 | } 75 | } 76 | 77 | ### Multi Query Example 78 | 79 | vector queries = vector{Query(query_string), select_query}; 80 | std::shared_ptr query_op = 81 | Connection::beginMultiQuery(std::move(conn), std::move(queries)); 82 | query_op->setCallback(query_callback); 83 | query_op->setTimeout(500ms); 84 | query_op->run(); 85 | 86 | void query_callback(MultiQueryOperation& op, 87 | QueryResult* query_result, 88 | QueryCallbackReason reason) { 89 | if (reason == QueryCallbackReason::RowsFetched) { 90 | LOG(INFO) << "Saw " << query_result->numRows() << " more rows!"; 91 | for (const auto& row : *query_result) { 92 | LOG(INFO) << "Row: " << folly::join(row, "\t"); 93 | } 94 | } else if (reason == QueryCallbackReason::QueryBoundary) { 95 | LOG(INFO) << "Finished statement " << query_result->queryNum(); 96 | } else if (reason == QueryCallbackReason::Success) { 97 | LOG(INFO) << "Query succeeded!"; 98 | } else { 99 | LOG(ERROR) << "Query failed! " << op.mysql_error(); 100 | } 101 | } 102 | 103 | ### Concepts 104 | 105 | * `AsyncMysqlClient` -- the main class that you use to get connections 106 | * `Connection` -- a connection to a specific database; pass this around 107 | when you run queries 108 | * `Operation`, `ConnectOperation`, `QueryOperation` `MultiQueryOperation` 109 | -- a virtual base class and its three concrete children that represent 110 | things you want to do (or have done) asynchronously 111 | * `OperationBatch` -- allows waiting for multiple operations to finish in 112 | parallel 113 | 114 | Asynchronous programming can be tricky. The Async MySQL client tried 115 | to manage this complexity via the concept of Operations. Operation 116 | objects represent some kind of async request, be it a connect or 117 | query, and hold the result when it completes. In the above example, 118 | we create a connection operation (`connect_op`), run it (which is a 119 | non-blocking call), and wait for it to complete (which is a blocking 120 | call). Specifically, *no function on any method of any class blocks, 121 | except for wait(). 122 | 123 | In addition to being able to start a MySQL query or connection "in the 124 | background" like the connect above, you can specify a callback, as 125 | seen in the query example. The callback may be invoked multiple 126 | times; for each callback it's passed the reason for it: RowsFetched, 127 | QueryBoundary when you're running multiple queries, Success and Failure. 128 | The data regarding the query is inside QueryResult (last insert id, 129 | affected rows, etc). 130 | 131 | ### Ownership and Thread Safety 132 | 133 | 99.9% of the time, you don't need to own an actual client object; 134 | simply use `AsyncMysqlClient::defaultClient()` to get a singleton. 135 | This class is thread safe. If you do need multiple clients, just 136 | construct them yourself. 137 | 138 | One difficulty of asynchronous operations is ownership of the 139 | underlying requests and results. For this async client, 140 | `std::shared_ptr` and `std::unique_ptr` are used to manage the 141 | relevant ownerships. Connections are owned via a `unique_ptr` and are 142 | handed over to the async client when a query is being run. When the 143 | query completes (either successfully or with any error except a 144 | timeout), you can retrieve the connection to begin a new query via 145 | `releaseConnection`. 146 | 147 | Operations are shared ownership with the client itself. Once you call 148 | `run()` however, you should not invoke other methods besides `wait()` 149 | as status checks, etc, will otherwise be racey. If you're using a 150 | callback, you can safely discard your reference to the operation; 151 | cleanup will occur when the query completes (either with rows or 152 | errors). 153 | 154 | ### Bonus feature: Query Formatting and Construction 155 | 156 | The mysql_client also supports a new method of constructing queries, 157 | similar to the www function `queryfx`. See `Query.h` for details, but 158 | in a nutshell, it is designed to prevent SQL injection while making 159 | query construction itself easier. 160 | -------------------------------------------------------------------------------- /squangle/mysql_client/ResetOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/ResetOperation.h" 10 | #include "squangle/mysql_client/Connection.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | InternalConnection::Status ResetOperation::runSpecialOperation() { 15 | return conn().getInternalConnection().resetConn(); 16 | } 17 | 18 | } // namespace facebook::common::mysql_client 19 | -------------------------------------------------------------------------------- /squangle/mysql_client/ResetOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/SpecialOperation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | // This is for sending COM_RESET_CONNECTION command before returning an idle 16 | // connection back to connection pool 17 | class ResetOperation : public SpecialOperation { 18 | public: 19 | explicit ResetOperation(std::unique_ptr impl) 20 | : SpecialOperation(std::move(impl)) {} 21 | 22 | private: 23 | InternalConnection::Status runSpecialOperation() override; 24 | 25 | db::OperationType getOperationType() const override { 26 | return db::OperationType::Reset; 27 | } 28 | 29 | const char* getErrorMsg() const override { 30 | return errorMsg; 31 | } 32 | 33 | static constexpr const char* errorMsg = "Reset connection failed: "; 34 | }; 35 | 36 | } // namespace facebook::common::mysql_client 37 | -------------------------------------------------------------------------------- /squangle/mysql_client/SSLOptionsProviderBase.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/SSLOptionsProviderBase.h" 10 | #include 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | bool SSLOptionsProviderBase::setMysqlSSLOptions(MYSQL* mysql) { 15 | auto sslContext = getSSLContext(); 16 | if (!sslContext) { 17 | return false; 18 | } 19 | // We need to set ssl_mode because we set it to disabled after we call 20 | // mysql_init. 21 | enum mysql_ssl_mode ssl_mode = SSL_MODE_PREFERRED; 22 | mysql_options(mysql, MYSQL_OPT_SSL_MODE, &ssl_mode); 23 | mysql_options(mysql, MYSQL_OPT_SSL_CONTEXT, sslContext->getSSLCtx()); 24 | auto sslSession = getRawSSLSession(); 25 | if (sslSession) { 26 | mysql_options4( 27 | mysql, MYSQL_OPT_SSL_SESSION, sslSession.release(), (void*)1); 28 | } 29 | return true; 30 | } 31 | 32 | bool SSLOptionsProviderBase::storeMysqlSSLSession(MYSQL* mysql) { 33 | auto reused = mysql_get_ssl_session_reused(mysql); 34 | if (!reused) { 35 | folly::ssl::SSLSessionUniquePtr session( 36 | (SSL_SESSION*)mysql_get_ssl_session(mysql)); 37 | if (session) { 38 | storeRawSSLSession(std::move(session)); 39 | } 40 | } 41 | return reused; 42 | } 43 | 44 | } // namespace facebook::common::mysql_client 45 | -------------------------------------------------------------------------------- /squangle/mysql_client/SSLOptionsProviderBase.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | 15 | namespace folly { 16 | class SSLContext; 17 | } 18 | 19 | namespace facebook::common::mysql_client { 20 | 21 | /* Interface for an SSL connection options Provider */ 22 | class SSLOptionsProviderBase { 23 | public: 24 | virtual ~SSLOptionsProviderBase() = default; 25 | 26 | // The SSL Context and Session options to be set for the connection 27 | virtual std::shared_ptr getSSLContext() = 0; 28 | 29 | // These sessions are raw OpenSSL sessions, currently used for resumption 30 | // in MySQL client 31 | virtual folly::ssl::SSLSessionUniquePtr getRawSSLSession() = 0; 32 | virtual void storeRawSSLSession( 33 | folly::ssl::SSLSessionUniquePtr ssl_session) = 0; 34 | 35 | // These sessions are abstracted ssl sessions, currently used for 36 | // resumption with folly::AsyncSSLSocket 37 | virtual std::shared_ptr getSSLSession() = 0; 38 | virtual void storeSSLSession( 39 | std::shared_ptr ssl_session) = 0; 40 | virtual void allowSessionResumption(bool allow) = 0; 41 | 42 | // Set the SSL Options on the MYSQL object. 43 | // Returns true if set was successful. 44 | bool setMysqlSSLOptions(MYSQL* mysql); 45 | 46 | // Fetches the SSL Session from the MYSQL object and stores it. 47 | // Returns if the SSL Session was reused for this connection. 48 | bool storeMysqlSSLSession(MYSQL* mysql); 49 | }; 50 | 51 | } // namespace facebook::common::mysql_client 52 | -------------------------------------------------------------------------------- /squangle/mysql_client/SemiFutureAdapter.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | 11 | #include "squangle/mysql_client/AsyncHelpers.h" 12 | #include "squangle/mysql_client/ConnectOperation.h" 13 | #include "squangle/mysql_client/Connection.h" 14 | #include "squangle/mysql_client/MultiQueryOperation.h" 15 | #include "squangle/mysql_client/Operation.h" 16 | #include "squangle/mysql_client/QueryOperation.h" 17 | #include "squangle/mysql_client/SemiFutureAdapter.h" 18 | 19 | namespace facebook::common::mysql_client { 20 | 21 | void handleConnectionCompletion( 22 | ConnectOperation& op, 23 | folly::Promise& promise) { 24 | auto conn = op.releaseConnection(); 25 | if (op.ok()) { 26 | // succeeded, let's build the result 27 | promise.setValue(ConnectResult( 28 | std::move(conn), 29 | op.result(), 30 | op.getKey(), 31 | op.opElapsed(), 32 | op.attemptsMade())); 33 | } else { 34 | // failed - build the exception 35 | promise.setException(conn->client().exceptionBuilder().buildMysqlException( 36 | op.result(), 37 | op.mysql_errno(), 38 | op.mysql_error(), 39 | conn->getKey(), 40 | op.opElapsed())); 41 | } 42 | } 43 | 44 | folly::SemiFuture toSemiFuture( 45 | const ConnectOperation_ptr& conn_op) { 46 | folly::MoveWrapper> promise; 47 | auto future = promise->getSemiFuture(); 48 | 49 | conn_op->setCallback([promise](ConnectOperation& op) mutable { 50 | handleConnectionCompletion(op, *promise); 51 | }); 52 | conn_op->run(); 53 | return future; 54 | } 55 | 56 | template 57 | void handleQueryCompletion( 58 | Operation& op, 59 | QueryResult query_result, 60 | QueryCallbackReason reason, 61 | folly::Promise>& promise) { 62 | auto conn = op.releaseConnection(); 63 | // Make a copy of the connection key here because we move `conn` below. 64 | auto conn_key = conn->getKey(); 65 | if (reason == QueryCallbackReason::Success) { 66 | ResultType result( 67 | std::move(query_result), 68 | op.numQueriesExecuted(), 69 | op.resultSize(), 70 | std::move(conn), 71 | op.result(), 72 | std::move(conn_key), 73 | op.opElapsed()); 74 | promise.setValue( 75 | std::make_pair(std::move(result), op.stealPostQueryCallback())); 76 | } else { 77 | promise.setException(conn->client().exceptionBuilder().buildQueryException( 78 | op.numQueriesExecuted(), 79 | op.result(), 80 | op.mysql_errno(), 81 | op.mysql_error(), 82 | std::move(conn_key), 83 | op.opElapsed())); 84 | } 85 | } 86 | 87 | // Handle pre-query callbacks (if they exist) 88 | template 89 | folly::SemiFuture handlePreQueryCallback(Operation& op) { 90 | // Use the pre-query callback if we have it, or else an empty SemiFuture 91 | auto optFut = op.callPreQueryCallback(op); 92 | if (optFut) { 93 | return std::move(*optFut); 94 | } 95 | 96 | return folly::makeSemiFuture(folly::unit); 97 | } 98 | 99 | // Handle setting up the promise, the callbacks when the query is finished and 100 | // running the query 101 | template 102 | folly::SemiFuture> handleRunQuery( 103 | const std::shared_ptr& op) { 104 | folly::MoveWrapper< 105 | folly::Promise>> 106 | promise; 107 | auto future = promise->getSemiFuture(); 108 | 109 | op->chainCallback(resultAppender([promise]( 110 | Operation& op, 111 | auto query_result, 112 | QueryCallbackReason reason) mutable { 113 | handleQueryCompletion( 114 | op, std::move(query_result), reason, *promise); 115 | })); 116 | op->run(); 117 | return future; 118 | } 119 | 120 | // Handle running the post-query callbacks (if they exist). The query result 121 | // will be in the first field of the result pair while the callback will be in 122 | // the second field of the pair. 123 | template 124 | folly::SemiFuture handlePostQueryCallback( 125 | std::pair&& resultPair) { 126 | if (resultPair.second) { 127 | // If we have a callback set, wrap (and then unwrap) the 128 | // result to/from the callback's std::variant wrapper 129 | return resultPair.second(AsyncPostQueryResult(std::move(resultPair.first))) 130 | .deferValue([](AsyncPostQueryResult&& result) { 131 | return std::get(std::move(result)); 132 | }); 133 | } 134 | return folly::makeSemiFuture(std::move(resultPair.first)); 135 | } 136 | 137 | // Convert the query operation into a semifuture, starting with the pre-query 138 | // callbacks then the query then the post-query callbacks. 139 | template 140 | folly::SemiFuture toSemiFutureHelper( 141 | std::shared_ptr op) { 142 | // Run pre-query callbacks 143 | auto sfut1 = handlePreQueryCallback(*op); 144 | auto sfut2 = std::move(sfut1).deferValue( 145 | [=](auto&& /* unused */) { return handleRunQuery(op); }); 146 | 147 | return std::move(sfut2).deferValue([op = std::move(op)](auto&& result) { 148 | // Pass `op` into this lambda to verify it's lifetime until we have handled 149 | // post query callbacks 150 | return handlePostQueryCallback(std::move(result)); 151 | }); 152 | } 153 | 154 | folly::SemiFuture toSemiFuture(QueryOperation_ptr query_op) { 155 | return toSemiFutureHelper( 156 | std::move(query_op)); 157 | } 158 | 159 | folly::SemiFuture toSemiFuture( 160 | MultiQueryOperation_ptr mquery_op) { 161 | return toSemiFutureHelper< 162 | DbMultiQueryResult, 163 | MultiQueryOperation, 164 | std::vector>(std::move(mquery_op)); 165 | } 166 | 167 | } // namespace facebook::common::mysql_client 168 | -------------------------------------------------------------------------------- /squangle/mysql_client/SemiFutureAdapter.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | /* 10 | * toFuture is our interface for database operation using folly::Future. 11 | * It's completely compatible with the `Operation` interface, so to use 12 | * futures all you need is to pass the `Operation`. 13 | * 14 | */ 15 | 16 | #pragma once 17 | 18 | #include "squangle/mysql_client/DbResult.h" 19 | 20 | #include 21 | 22 | namespace facebook::common::mysql_client { 23 | 24 | class ConnectOperation; 25 | class QueryOperation; 26 | class MultiQueryOperation; 27 | 28 | typedef std::shared_ptr ConnectOperation_ptr; 29 | typedef std::shared_ptr QueryOperation_ptr; 30 | typedef std::shared_ptr MultiQueryOperation_ptr; 31 | 32 | // SemiFuture for ConnectOperation 33 | FOLLY_NODISCARD folly::SemiFuture toSemiFuture( 34 | const ConnectOperation_ptr& conn_op); 35 | 36 | // SemiFuture for QueryOperation 37 | FOLLY_NODISCARD folly::SemiFuture toSemiFuture( 38 | QueryOperation_ptr query_op); 39 | 40 | // SemiFuture for MultiQueryOperation 41 | FOLLY_NODISCARD folly::SemiFuture toSemiFuture( 42 | MultiQueryOperation_ptr mquery_op); 43 | 44 | } // namespace facebook::common::mysql_client 45 | -------------------------------------------------------------------------------- /squangle/mysql_client/SpecialOperation.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/SpecialOperation.h" 10 | #include "squangle/mysql_client/Connection.h" 11 | 12 | namespace facebook::common::mysql_client { 13 | 14 | SpecialOperation& SpecialOperationImpl::getOp() const { 15 | DCHECK(op_ && dynamic_cast(op_) != nullptr); 16 | return *(SpecialOperation*)op_; 17 | } 18 | 19 | InternalConnection::Status SpecialOperationImpl::runSpecialOperation() { 20 | return getOp().runSpecialOperation(); 21 | } 22 | 23 | } // namespace facebook::common::mysql_client 24 | -------------------------------------------------------------------------------- /squangle/mysql_client/SpecialOperation.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/Operation.h" 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | class ConnectionProxy; 16 | class SpecialOperation; 17 | 18 | using SpecialOperationCallback = 19 | std::function; 20 | 21 | class SpecialOperationImpl : virtual public OperationBase { 22 | public: 23 | // SpecialOperationImpl() : OperationBase(nullptr) {} 24 | virtual ~SpecialOperationImpl() override = default; 25 | 26 | void setCallback(SpecialOperationCallback callback) { 27 | callback_ = std::move(callback); 28 | } 29 | 30 | InternalConnection::Status runSpecialOperation(); 31 | 32 | protected: 33 | [[nodiscard]] SpecialOperation& getOp() const; 34 | 35 | SpecialOperationCallback callback_{nullptr}; 36 | }; 37 | 38 | // SpecialOperation means operations like COM_RESET_CONNECTION, 39 | // COM_CHANGE_USER, etc. 40 | class SpecialOperation : public Operation { 41 | public: 42 | void setCallback(SpecialOperationCallback callback) { 43 | impl_->setCallback(std::move(callback)); 44 | } 45 | 46 | protected: 47 | explicit SpecialOperation(std::unique_ptr impl) 48 | : impl_(std::move(impl)) { 49 | if (!impl_) { 50 | throw std::runtime_error("ConnectOperationImpl is null"); 51 | } 52 | 53 | impl_->setOperation(*this); 54 | } 55 | 56 | virtual InternalConnection::Status runSpecialOperation() = 0; 57 | friend SpecialOperationImpl; 58 | 59 | private: 60 | virtual const char* getErrorMsg() const = 0; 61 | 62 | virtual OperationBase* impl() override { 63 | return (OperationBase*)impl_.get(); 64 | } 65 | virtual const OperationBase* impl() const override { 66 | return (OperationBase*)impl_.get(); 67 | } 68 | 69 | std::unique_ptr impl_; 70 | }; 71 | 72 | } // namespace facebook::common::mysql_client 73 | -------------------------------------------------------------------------------- /squangle/mysql_client/SyncConnectionPool.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/SyncConnectionPool.h" 10 | 11 | using namespace std::chrono_literals; 12 | 13 | namespace facebook::common::mysql_client { 14 | 15 | std::shared_ptr SyncConnectionPool::makePool( 16 | std::shared_ptr mysql_client, 17 | PoolOptions pool_options) { 18 | auto connectionPool = std::make_shared( 19 | std::move(mysql_client), std::move(pool_options)); 20 | return connectionPool; 21 | } 22 | 23 | std::unique_ptr SyncConnectionPool::connect( 24 | const std::string& host, 25 | int port, 26 | const std::string& database_name, 27 | const std::string& user, 28 | const std::string& password, 29 | const ConnectionOptions& conn_opts) { 30 | auto op = beginConnection(host, port, database_name, user, password); 31 | op->setConnectionOptions(conn_opts); 32 | // This will throw (intended behaviour) in case the operation didn't succeed 33 | return blockingConnectHelper(*op); 34 | } 35 | 36 | void SyncConnectionPool::openNewConnectionPrep( 37 | SyncConnectPoolOperation& pool_op) { 38 | pool_op.prepWait(); 39 | } 40 | 41 | void SyncConnectionPool::openNewConnectionFinish( 42 | SyncConnectPoolOperation& pool_op, 43 | const PoolKey& pool_key) { 44 | if (!pool_op.syncWait()) { 45 | if (!conn_storage_.dequeueOperation(pool_key, pool_op)) { 46 | // The operation was not found in the queue, so someone must be fulfilling 47 | // the operation. Wait until that is finished. 48 | while (pool_op.isActive()) { 49 | /* sleep_override */ std::this_thread::sleep_for(1ms); 50 | } 51 | 52 | return; 53 | } 54 | 55 | pool_op.timeoutTriggered(); 56 | } 57 | 58 | pool_op.cleanupWait(); 59 | } 60 | 61 | template <> 62 | void SyncConnectPoolOperationImpl::specializedRun() { 63 | // No special thread manipulation needed for sync client 64 | MysqlConnectPoolOperationImpl::specializedRunImpl(); 65 | } 66 | 67 | template <> 68 | std::unique_ptr> 69 | createConnectPoolOperationImpl( 70 | std::weak_ptr> pool, 71 | std::shared_ptr client, 72 | std::shared_ptr conn_key) { 73 | return std::make_unique< 74 | mysql_protocol::MysqlConnectPoolOperationImpl>( 75 | std::move(pool), client, std::move(conn_key)); 76 | } 77 | 78 | } // namespace facebook::common::mysql_client 79 | -------------------------------------------------------------------------------- /squangle/mysql_client/SyncConnectionPool.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | 13 | #include "squangle/mysql_client/ConnectPoolOperation.h" 14 | #include "squangle/mysql_client/ConnectionPool.h" 15 | #include "squangle/mysql_client/SyncMysqlClient.h" 16 | #include "squangle/mysql_client/mysql_protocol/MysqlConnectPoolOperationImpl.h" 17 | 18 | namespace facebook::common::mysql_client { 19 | 20 | using SyncConnectPoolOperation = ConnectPoolOperation; 21 | using SyncConnectPoolOperationImpl = 22 | mysql_protocol::MysqlConnectPoolOperationImpl; 23 | 24 | class SyncConnectionPool : public ConnectionPool { 25 | public: 26 | // Don't use std::chrono::duration::MAX to avoid overflows 27 | static std::shared_ptr makePool( 28 | std::shared_ptr mysql_client, 29 | PoolOptions pool_options = PoolOptions()); 30 | 31 | std::unique_ptr connect( 32 | const std::string& host, 33 | int port, 34 | const std::string& database_name, 35 | const std::string& user, 36 | const std::string& password, 37 | const ConnectionOptions& conn_opts = ConnectionOptions()); 38 | 39 | // Don't use the constructor directly, only public to use make_shared 40 | SyncConnectionPool( 41 | std::shared_ptr mysql_client, 42 | PoolOptions pool_options) 43 | : ConnectionPool( 44 | std::move(mysql_client), 45 | std::move(pool_options)) { 46 | scheduler_.addFunction( 47 | [&]() { 48 | conn_storage_.cleanupOperations(); 49 | conn_storage_.cleanupConnections(); 50 | }, 51 | PoolOptions::kCleanUpTimeout, 52 | "pool_periodic_cleanup"); 53 | scheduler_.start(); 54 | } 55 | 56 | ~SyncConnectionPool() override { 57 | VLOG(2) << "Connection pool dying"; 58 | 59 | shutdown(); 60 | 61 | VLOG(2) << "Connection pool shutdown completed"; 62 | } 63 | 64 | SyncConnectionPool(const SyncConnectionPool&) = delete; 65 | SyncConnectionPool& operator=(const SyncConnectionPool&) = delete; 66 | 67 | SyncConnectionPool(SyncConnectionPool&&) = delete; 68 | SyncConnectionPool& operator=(SyncConnectionPool&&) = delete; 69 | 70 | void shutdown() override { 71 | bool expected = false; 72 | if (shutting_down_.compare_exchange_strong(expected, true)) { 73 | scheduler_.shutdown(); 74 | conn_storage_.clearAll(); 75 | } 76 | } 77 | 78 | private: 79 | bool isShuttingDown() const override { 80 | return shutting_down_; 81 | } 82 | 83 | void validateCorrectThread() const override { 84 | // The sync connection pool runs everything in the clients' threads so don't 85 | // do anything here. 86 | } 87 | 88 | bool runInCorrectThread(std::function&& func, bool /*wait*/) 89 | override { 90 | // The sync connection pool runs everything in the clients' threads. 91 | func(); 92 | return true; 93 | } 94 | 95 | std::unique_ptr makeNewConnection( 96 | std::shared_ptr conn_key, 97 | std::unique_ptr> mysqlConn) override { 98 | return std::make_unique( 99 | *mysql_client_, std::move(conn_key), std::move(mysqlConn)); 100 | } 101 | 102 | void openNewConnectionPrep(SyncConnectPoolOperation& pool_op) override; 103 | 104 | void openNewConnectionFinish( 105 | SyncConnectPoolOperation& pool_op, 106 | const PoolKey& pool_key) override; 107 | 108 | std::atomic shutting_down_{false}; 109 | 110 | folly::FunctionScheduler scheduler_; 111 | }; 112 | 113 | } // namespace facebook::common::mysql_client 114 | -------------------------------------------------------------------------------- /squangle/mysql_client/SyncMysqlClient.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | #include "squangle/mysql_client/ResetOperation.h" 13 | #include "squangle/mysql_client/SyncMysqlClient.h" 14 | #include "squangle/mysql_client/mysql_protocol/MysqlConnectOperationImpl.h" 15 | #include "squangle/mysql_client/mysql_protocol/MysqlFetchOperationImpl.h" 16 | #include "squangle/mysql_client/mysql_protocol/MysqlSpecialOperationImpl.h" 17 | 18 | namespace facebook::common::mysql_client { 19 | 20 | namespace detail { 21 | 22 | struct SyncMysqlClientSingletonTag {}; 23 | 24 | folly::Singleton 25 | defaultSyncMysqlClientSingleton; 26 | 27 | } // namespace detail 28 | 29 | std::shared_ptr SyncMysqlClient::defaultClient() { 30 | return detail::defaultSyncMysqlClientSingleton.try_get(); 31 | } 32 | 33 | std::unique_ptr 34 | SyncMysqlClient::createConnectOperationImpl( 35 | MysqlClientBase* client_base, 36 | std::shared_ptr conn_key) const { 37 | return std::make_unique( 38 | client_base, std::move(conn_key)); 39 | } 40 | 41 | std::unique_ptr SyncMysqlClient::createFetchOperationImpl( 42 | std::unique_ptr conn) const { 43 | return std::make_unique( 44 | std::move(conn)); 45 | } 46 | 47 | std::unique_ptr 48 | SyncMysqlClient::createSpecialOperationImpl( 49 | std::unique_ptr conn) const { 50 | return std::make_unique( 51 | std::move(conn)); 52 | } 53 | 54 | std::unique_ptr SyncMysqlClient::createConnection( 55 | std::shared_ptr conn_key) { 56 | return std::make_unique(*this, std::move(conn_key)); 57 | } 58 | 59 | SyncConnection::~SyncConnection() { 60 | if (mysql_connection_ && conn_dying_callback_ && needToCloneConnection_ && 61 | isReusable() && !inTransaction() && 62 | getConnectionOptions().isEnableResetConnBeforeClose()) { 63 | // We clone this Connection object to send COM_RESET_CONNECTION command 64 | // via the connection before returning it to the connection pool. 65 | // The callback function points to recycleMysqlConnection(), which is 66 | // responsible for recyclining the connection. 67 | // This object's callback is set to null and the cloned object's 68 | // callback instead points to the original callback function, which will 69 | // be called after COM_RESET_CONNECTION. 70 | 71 | auto connHolder = stealConnectionHolder(true); 72 | auto conn = std::make_unique( 73 | client(), getKey(), std::move(connHolder)); 74 | conn->needToCloneConnection_ = false; 75 | conn->setConnectionOptions(getConnectionOptions()); 76 | conn->setConnectionDyingCallback(std::move(conn_dying_callback_)); 77 | conn_dying_callback_ = nullptr; 78 | auto resetOp = Connection::resetConn(std::move(conn)); 79 | resetOp->run().wait(); 80 | } 81 | } 82 | 83 | } // namespace facebook::common::mysql_client 84 | -------------------------------------------------------------------------------- /squangle/mysql_client/SyncMysqlClient.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/Connection.h" 12 | #include "squangle/mysql_client/MysqlClientBase.h" 13 | #include "squangle/mysql_client/mysql_protocol/MysqlConnection.h" 14 | 15 | namespace facebook::common::mysql_client { 16 | 17 | class SyncConnection; 18 | 19 | class SyncMysqlClient : public MysqlClientBase { 20 | public: 21 | SyncMysqlClient() : SyncMysqlClient(nullptr) {} 22 | explicit SyncMysqlClient( 23 | std::unique_ptr db_logger, 24 | std::unique_ptr db_stats = 25 | std::make_unique(), 26 | std::unique_ptr exception_builder = nullptr) 27 | : MysqlClientBase( 28 | adjustLogger(std::move(db_logger)), 29 | std::move(db_stats), 30 | std::move(exception_builder)) {} 31 | 32 | static std::unique_ptr adjustLogger( 33 | std::unique_ptr logger) { 34 | if (logger) { 35 | logger->setLoggingPrefix("cpp_sync"); 36 | } 37 | return logger; 38 | } 39 | 40 | // Factory method 41 | std::unique_ptr createConnection( 42 | std::shared_ptr conn_key) override; 43 | 44 | void drain(bool /*unused*/) {} 45 | 46 | bool supportsLocalFiles() override { 47 | return true; 48 | } 49 | 50 | uint64_t getPoolsConnectionLimit() { 51 | // This is used by HHVM in the async client. We don't need it here in the 52 | // sync client. 53 | return std::numeric_limits::max(); 54 | } 55 | 56 | static std::shared_ptr defaultClient(); 57 | 58 | protected: 59 | // Private methods, primarily used by Operations and its subclasses. 60 | template 61 | friend class ConnectionPool; 62 | 63 | std::unique_ptr createConnectOperationImpl( 64 | MysqlClientBase* client, 65 | std::shared_ptr conn_key) const override; 66 | 67 | std::unique_ptr createFetchOperationImpl( 68 | std::unique_ptr conn) const override; 69 | 70 | std::unique_ptr createSpecialOperationImpl( 71 | std::unique_ptr conn) const override; 72 | }; 73 | 74 | // SyncConnection is a specialization of Connection to handle inline loops. 75 | // It has its own EventBase and Operations using it will have events and 76 | // callbacks running with this EventBase. 77 | class SyncConnection : public Connection { 78 | public: 79 | SyncConnection( 80 | MysqlClientBase& client, 81 | std::shared_ptr conn_key, 82 | std::unique_ptr conn = nullptr) 83 | : Connection(client, std::move(conn_key), std::move(conn)) {} 84 | 85 | ~SyncConnection() override; 86 | 87 | SyncConnection(const SyncConnection&) = delete; 88 | SyncConnection& operator=(const SyncConnection&) = delete; 89 | 90 | SyncConnection(SyncConnection&&) = delete; 91 | SyncConnection& operator=(SyncConnection&&) = delete; 92 | 93 | std::shared_ptr createOperation( 94 | std::unique_ptr proxy, 95 | MultiQuery&& multi_query) override { 96 | auto impl = client().createFetchOperationImpl(std::move(proxy)); 97 | return MultiQueryStreamOperation::create( 98 | std::move(impl), std::move(multi_query)); 99 | } 100 | 101 | protected: 102 | std::unique_ptr createInternalConnection() override { 103 | return std::make_unique(); 104 | } 105 | }; 106 | 107 | } // namespace facebook::common::mysql_client 108 | -------------------------------------------------------------------------------- /squangle/mysql_client/TODO: -------------------------------------------------------------------------------- 1 | consistency of column vs field terminology 2 | 3 | Support multiquery 4 | Connection pooling 5 | Logging/monitoring 6 | Retry? Reconnect? 7 | 8 | Support use of an external TEventBase 9 | -------------------------------------------------------------------------------- /squangle/mysql_client/TwoLevelCache.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace { 17 | 18 | // Find the best match based on the results of pred. `pred` returns a 19 | // std::optional, and returning std::nullopt indicates that the loop 20 | // should be terminated an the current best returned. Otherwise true indicates 21 | // that the current value (`*it`) is better than `*best` and false means it is 22 | // not. 23 | template 24 | auto findBest(Container& container, Compare comp) { 25 | auto it = container.begin(); 26 | if (it == container.end()) { 27 | return container.end(); 28 | } 29 | 30 | auto best = it; 31 | while (++it != container.end()) { 32 | if (auto res = comp(*best, *it); !res) { 33 | break; 34 | } else if (*res) { 35 | best = it; 36 | } 37 | } 38 | 39 | return best; 40 | } 41 | 42 | } // namespace 43 | 44 | namespace facebook::common::mysql_client { 45 | 46 | // level1_ map: Key = PoolKey, Value = list of MysqlPooledHolder 47 | // level2_ map: Key = PoolKey (dbname is ignored), Value = set of PoolKey 48 | // 49 | // Pool keys will be present in level2_ set as long as its level1_ list include 50 | // non-zero connections. 51 | template < 52 | typename Client, 53 | typename Key, 54 | typename Value, 55 | typename FullKeyHash, 56 | typename PartialKeyHash> 57 | class TwoLevelCache { 58 | public: 59 | TwoLevelCache() {} 60 | 61 | void push(const Key& key, Value value, size_t max) { 62 | auto& list = level1_[key]; 63 | if (list.empty()) { 64 | // The key is new in level1_, let's add it to level2_. 65 | level2_[key].insert(key); 66 | } 67 | 68 | list.push_back(std::move(value)); 69 | if (list.size() > max) { 70 | list.pop_front(); 71 | } 72 | } 73 | 74 | Value popLevel1(const Key& key) { 75 | if (auto it = level1_.find(key); 76 | it != level1_.end() && it->second.size() > 0) { 77 | auto ret = std::move(it->second.front()); 78 | it->second.pop_front(); 79 | 80 | if (it->second.empty()) { 81 | // The key does not exist in level1_, let's remove it from level2_. 82 | eraseFromLevel2(key); 83 | level1_.erase(it); 84 | } 85 | return ret; 86 | } 87 | return Value(); 88 | } 89 | 90 | // Finds a key from level2_ cache, with a predicate comp, and pops its 91 | // corresponding connection from level1_ cache. 92 | // Predicate comp returns true, if the best key needs to be updated, false, if 93 | // the best key remains unchanged, or folly::none, if we need to stop 94 | // traversing keys in level2_. 95 | template 96 | Value popLevel2(const Key& key, Pred comp) { 97 | if (auto it = level2_.find(key); it != level2_.end()) { 98 | auto best = findBest(it->second, std::move(comp)); 99 | DCHECK(best != it->second.end()); 100 | return popLevel1(*best); 101 | } 102 | return Value(); 103 | } 104 | 105 | // Cleans up connections in level2_ (and level1_) cache, which meet the pred. 106 | template 107 | std::vector cleanup(Pred pred) { 108 | std::vector toBeReleased; 109 | for (auto it1 = level1_.begin(); it1 != level1_.end();) { 110 | auto& list = it1->second; 111 | DCHECK(!list.empty()); 112 | 113 | for (auto it2 = list.begin(); it2 != list.end();) { 114 | if (pred(*it2)) { 115 | // Add the value to the toBeReleased vector so all of them get 116 | // destructed after the loop completes. 117 | toBeReleased.push_back(std::move(*it2)); 118 | it2 = list.erase(it2); 119 | } else { 120 | ++it2; 121 | } 122 | } 123 | 124 | if (list.empty()) { 125 | // The key does not exist in level1_, let's remove it from level2_ 126 | eraseFromLevel2(it1->first); 127 | it1 = level1_.erase(it1); 128 | } else { 129 | ++it1; 130 | } 131 | } 132 | 133 | return toBeReleased; 134 | } 135 | 136 | void clear() { 137 | level1_.clear(); 138 | level2_.clear(); 139 | } 140 | 141 | template 142 | void iterateLevel1(Func func) { 143 | for (const auto& [key, value] : level1_) { 144 | func(key, value); 145 | } 146 | } 147 | 148 | template 149 | void iterateLevel2(Func func) { 150 | for (const auto& [key, value] : level2_) { 151 | func(key, value); 152 | } 153 | } 154 | 155 | size_t level1Size(const Key& key) const { 156 | if (auto it = level1_.find(key); it != level1_.end()) { 157 | return it->second.size(); 158 | } 159 | return 0; 160 | } 161 | 162 | size_t level2Size(const Key& key) const { 163 | if (auto it = level2_.find(key); it != level2_.end()) { 164 | return it->second.size(); 165 | } 166 | return 0; 167 | } 168 | 169 | size_t level1NumKey() const noexcept { 170 | return level1_.size(); 171 | } 172 | 173 | private: 174 | using Level2Value = folly::F14FastSet; 175 | using Level1Map = folly::F14FastMap, FullKeyHash>; 176 | using Level2Map = 177 | folly::F14FastMap; 178 | 179 | // This should be called before erasing the empty list from level1 180 | void eraseFromLevel2(const Key& key) { 181 | DCHECK(level1_.contains(key)); 182 | DCHECK(level1_.at(key).empty()); 183 | auto it = level2_.find(key); 184 | DCHECK(it != level2_.end()); 185 | it->second.erase(key); 186 | if (it->second.empty()) { 187 | level2_.erase(it); 188 | } 189 | } 190 | 191 | Level1Map level1_; 192 | Level2Map level2_; 193 | }; 194 | 195 | } // namespace facebook::common::mysql_client 196 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlConnectOperationImpl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/ConnectOperation.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlOperationImpl.h" 13 | 14 | namespace facebook::common::mysql_client::mysql_protocol { 15 | 16 | class MysqlConnectOperationImpl : public MysqlOperationImpl, 17 | virtual public ConnectOperationImpl { 18 | public: 19 | // Don't call this; it's public strictly for AsyncMysqlClient to be 20 | // able to call make_shared. 21 | MysqlConnectOperationImpl( 22 | MysqlClientBase* mysql_client, 23 | std::shared_ptr conn_key); 24 | virtual ~MysqlConnectOperationImpl() override; 25 | 26 | // copy and move not allowed 27 | MysqlConnectOperationImpl(const MysqlConnectOperationImpl&) = delete; 28 | MysqlConnectOperationImpl& operator=(const MysqlConnectOperationImpl&) = 29 | delete; 30 | 31 | MysqlConnectOperationImpl(MysqlConnectOperationImpl&&) = delete; 32 | MysqlConnectOperationImpl& operator=(MysqlConnectOperationImpl&&) = delete; 33 | 34 | static constexpr Duration kMinimumViableConnectTimeout = 35 | std::chrono::microseconds(50); 36 | 37 | bool isActive() const override { 38 | return active_in_client_; 39 | } 40 | 41 | protected: 42 | virtual void attemptFailed(OperationResult result); 43 | virtual void attemptSucceeded(OperationResult result); 44 | 45 | virtual void specializedRun() override; 46 | void actionable() override; 47 | void specializedTimeoutTriggered() override; 48 | void specializedCompleteOperation() override; 49 | 50 | // Called when tcp timeout is triggered 51 | void tcpConnectTimeoutTriggered(); 52 | 53 | // Removes the Client ref, it can be called by child classes without needing 54 | // to add them as friend classes of AsyncMysqlClient 55 | virtual void removeClientReference(); 56 | 57 | bool shouldCompleteOperation(OperationResult result); 58 | 59 | folly::ssl::SSLSessionUniquePtr getSSLSession(); 60 | 61 | // Implementation of timeout handling for tcpTimeout and overall connect 62 | // timeout 63 | void timeoutHandler(bool isTcpTimeout, bool isPool = false); 64 | 65 | private: 66 | virtual void specializedRunImpl(); 67 | 68 | void logConnectCompleted(OperationResult result); 69 | 70 | void maybeStoreSSLSession(); 71 | 72 | bool isDoneWithTcpHandShake(); 73 | 74 | const MysqlConnectionKey& getMysqlKeyRef() const { 75 | DCHECK(dynamic_cast(conn_key_.get())); 76 | return *((const MysqlConnectionKey*)conn_key_.get()); 77 | } 78 | 79 | static int mysqlCertValidator( 80 | X509* server_cert, 81 | const void* context, 82 | const char** errptr); 83 | 84 | int flags_; 85 | 86 | bool active_in_client_; 87 | 88 | // Timeout used for controlling early timeout of just the tcp handshake phase 89 | // before doing heavy lifting like ssl and other mysql protocol for connection 90 | // establishment 91 | class ConnectTcpTimeoutHandler : public folly::AsyncTimeout { 92 | public: 93 | ConnectTcpTimeoutHandler( 94 | folly::EventBase* base, 95 | MysqlConnectOperationImpl* connect_operation) 96 | : folly::AsyncTimeout(base), op_(connect_operation) {} 97 | 98 | ConnectTcpTimeoutHandler() = delete; 99 | ~ConnectTcpTimeoutHandler() override = default; 100 | 101 | // copy and move not allowed 102 | ConnectTcpTimeoutHandler(const ConnectTcpTimeoutHandler&) = delete; 103 | ConnectTcpTimeoutHandler& operator=(const ConnectTcpTimeoutHandler&) = 104 | delete; 105 | 106 | ConnectTcpTimeoutHandler(ConnectTcpTimeoutHandler&&) = delete; 107 | ConnectTcpTimeoutHandler& operator=(ConnectTcpTimeoutHandler&&) = delete; 108 | 109 | void timeoutExpired() noexcept override { 110 | op_->tcpConnectTimeoutTriggered(); 111 | } 112 | 113 | private: 114 | MysqlConnectOperationImpl* op_; 115 | }; 116 | 117 | ConnectTcpTimeoutHandler tcp_timeout_handler_; 118 | 119 | friend class AsyncMysqlClient; 120 | friend class MysqlClientBase; 121 | }; 122 | 123 | } // namespace facebook::common::mysql_client::mysql_protocol 124 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlFetchOperationImpl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/FetchOperation.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlOperationImpl.h" 13 | 14 | namespace facebook::common::mysql_client::mysql_protocol { 15 | 16 | class MysqlFetchOperationImpl : public MysqlOperationImpl, 17 | public FetchOperationImpl { 18 | public: 19 | explicit MysqlFetchOperationImpl( 20 | std::unique_ptr conn) 21 | : OperationBase(std::move(conn)) {} 22 | 23 | protected: 24 | void specializedRunImpl(); 25 | void specializedRun() override; 26 | 27 | // In actionable it is analyzed the action that is required to continue the 28 | // operation. For example, if the fetch action is StartQuery, it runs query or 29 | // requests more results depending if it had already ran or not the query. The 30 | // same process happens for the other FetchActions. The action member can be 31 | // changed in other member functions called in actionable to keep the fetching 32 | // flow running. 33 | void actionable() override; 34 | void specializedTimeoutTriggered() override; 35 | void specializedCompleteOperation() override; 36 | 37 | void cancelOp() { 38 | cancel_ = true; 39 | setFetchAction(FetchAction::CompleteQuery); 40 | } 41 | 42 | bool hasDataInNativeFormat() const override { 43 | return false; 44 | } 45 | 46 | void pauseForConsumer() override; 47 | void resume() override; 48 | bool isPaused() const override; 49 | 50 | private: 51 | void resumeImpl(); 52 | // Checks if the current thread has access to stream, or result data. 53 | bool isStreamAccessAllowed() const override; 54 | 55 | // Asynchronously kill a currently running query, returns 56 | // before the query is killed 57 | void killRunningQuery(); 58 | 59 | bool cancel_ = false; 60 | std::atomic resume_scheduled_{false}; 61 | }; 62 | 63 | } // namespace facebook::common::mysql_client::mysql_protocol 64 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlOperationImpl.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include 10 | 11 | #include "squangle/mysql_client/Connection.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlConnection.h" 13 | #include "squangle/mysql_client/mysql_protocol/MysqlOperationImpl.h" 14 | 15 | namespace facebook::common::mysql_client::mysql_protocol { 16 | 17 | void MysqlOperationImpl::protocolCompleteOperation(OperationResult result) { 18 | conn().runInThread(this, &MysqlOperationImpl::completeOperation, result); 19 | } 20 | 21 | MysqlOperationImpl::MysqlOperationImpl() 22 | : OperationBase(nullptr), 23 | EventHandler(client_.getEventBase()), 24 | AsyncTimeout(client_.getEventBase()) {} 25 | 26 | bool MysqlOperationImpl::isInEventBaseThread() const { 27 | return conn().isInEventBaseThread(); 28 | } 29 | 30 | bool MysqlOperationImpl::isEventBaseSet() const { 31 | return conn().getEventBase() != nullptr; 32 | } 33 | 34 | void MysqlOperationImpl::invokeActionable() { 35 | DCHECK(isInEventBaseThread()); 36 | auto guard = makeRequestGuard(); 37 | actionable(); 38 | } 39 | 40 | /*static*/ 41 | MysqlConnection* MysqlOperationImpl::getMysqlConnection( 42 | InternalConnection* conn) { 43 | // We don't want to pay the cost of doing a dynamic_cast in prod when this 44 | // should _always_ be a MysqlConnection. Thus just validate on debug builds. 45 | DCHECK(dynamic_cast(conn)); 46 | return static_cast(conn); 47 | } 48 | 49 | /*static*/ 50 | const MysqlConnection* MysqlOperationImpl::getMysqlConnection( 51 | const InternalConnection* conn) { 52 | // We don't want to pay the cost of doing a dynamic_cast in prod when this 53 | // should _always_ be a MysqlConnection. Thus just validate on debug builds. 54 | DCHECK(dynamic_cast(conn)); 55 | return static_cast(conn); 56 | } 57 | 58 | const MysqlConnection* MysqlOperationImpl::getMysqlConnection() const { 59 | return getMysqlConnection(&getInternalConnection()); 60 | } 61 | 62 | MysqlConnection* MysqlOperationImpl::getMysqlConnection() { 63 | return getMysqlConnection(&getInternalConnection()); 64 | } 65 | 66 | void MysqlOperationImpl::waitForActionable() { 67 | DCHECK(isInEventBaseThread()); 68 | 69 | auto event_mask = getMysqlConnection()->getReadWriteState(); 70 | 71 | if (hasOpElapsed(getTimeout())) { 72 | timeoutTriggered(); 73 | return; 74 | } 75 | 76 | auto leftUs = getTimeout() - opElapsed(); 77 | auto leftMs = std::chrono::duration_cast(leftUs); 78 | scheduleTimeout(leftMs.count()); 79 | registerHandler(event_mask); 80 | } 81 | 82 | void MysqlOperationImpl::handlerReady(uint16_t /*events*/) noexcept { 83 | DCHECK(conn().isInEventBaseThread()); 84 | CHECK_THROW( 85 | state() != OperationState::Completed && 86 | state() != OperationState::Unstarted, 87 | db::OperationStateException); 88 | 89 | if (state() == OperationState::Cancelling) { 90 | cancel(); 91 | } else { 92 | invokeActionable(); 93 | } 94 | } 95 | 96 | void MysqlOperationImpl::timeoutTriggered() { 97 | specializedTimeoutTriggered(); 98 | } 99 | 100 | void MysqlOperationImpl::completeOperation(OperationResult result) { 101 | DCHECK(isInEventBaseThread()); 102 | if (state() == OperationState::Completed) { 103 | return; 104 | } 105 | 106 | CHECK_THROW( 107 | state() == OperationState::Pending || 108 | state() == OperationState::Cancelling || 109 | state() == OperationState::Unstarted, 110 | db::OperationStateException); 111 | completeOperationInner(result); 112 | } 113 | 114 | void MysqlOperationImpl::completeOperationInner(OperationResult result) { 115 | setState(OperationState::Completed); 116 | setResult(result); 117 | setDuration(); 118 | if ((result == OperationResult::Cancelled || 119 | result == OperationResult::TimedOut) && 120 | conn().hasInitialized()) { 121 | // Cancelled/timed out ops leave our connection in an undefined 122 | // state. Close it to prevent trouble. 123 | conn().close(); 124 | } 125 | 126 | unregisterHandler(); 127 | cancelTimeout(); 128 | 129 | if (callbacks_.post_operation_callback_) { 130 | callbacks_.post_operation_callback_(*op_); 131 | } 132 | 133 | specializedCompleteOperation(); 134 | 135 | // call observer callback 136 | if (observer_callback_) { 137 | observer_callback_(*op_); 138 | } 139 | 140 | deferRemoveOperation(op_); 141 | } 142 | 143 | /*static*/ std::string MysqlOperationImpl::connectStageString( 144 | connect_stage stage) { 145 | return MysqlConnection::findConnectStageName(stage).value_or(""); 146 | } 147 | 148 | } // namespace facebook::common::mysql_client::mysql_protocol 149 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlOperationImpl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/Operation.h" 12 | 13 | namespace facebook::common::mysql_client::mysql_protocol { 14 | 15 | class MysqlConnection; 16 | 17 | // The abstract base for our available Operations. Subclasses share 18 | // intimate knowledge with the Operation class (most member variables 19 | // are protected). 20 | class MysqlOperationImpl : virtual public OperationBase, 21 | public folly::EventHandler, 22 | public folly::AsyncTimeout { 23 | public: 24 | // No public constructor. 25 | virtual ~MysqlOperationImpl() override = default; 26 | 27 | // copy and move not allowed 28 | MysqlOperationImpl(const MysqlOperationImpl&) = delete; 29 | MysqlOperationImpl& operator=(const MysqlOperationImpl&) = delete; 30 | 31 | MysqlOperationImpl(MysqlOperationImpl&&) = delete; 32 | MysqlOperationImpl& operator=(MysqlOperationImpl&&) = delete; 33 | 34 | Duration getMaxThreadBlockTime() { 35 | return max_thread_block_time_; 36 | } 37 | 38 | Duration getTotalThreadBlockTime() { 39 | return total_thread_block_time_; 40 | } 41 | 42 | void logThreadBlockTime(const folly::stop_watch sw) { 43 | auto block_time = sw.elapsed(); 44 | max_thread_block_time_ = std::max(max_thread_block_time_, block_time); 45 | total_thread_block_time_ += block_time; 46 | } 47 | 48 | static std::string connectStageString(connect_stage stage); 49 | 50 | protected: 51 | MysqlOperationImpl(); 52 | 53 | MysqlConnection* getMysqlConnection(); 54 | const MysqlConnection* getMysqlConnection() const; 55 | 56 | static MysqlConnection* getMysqlConnection(InternalConnection* conn); 57 | static const MysqlConnection* getMysqlConnection( 58 | const InternalConnection* conn); 59 | 60 | // Called when an Operation needs to wait for the data to be readable or 61 | // writable (aka actionable). 62 | void waitForActionable(); 63 | 64 | // Overridden in child classes and invoked when the status is actionable. This 65 | // function should either completeOperation or waitForActionable. 66 | virtual void actionable() = 0; 67 | 68 | // EventHandler override 69 | void handlerReady(uint16_t /*events*/) noexcept override; 70 | 71 | // AsyncTimeout override 72 | void timeoutExpired() noexcept override { 73 | timeoutTriggered(); 74 | } 75 | 76 | // Called by AsyncTimeout::timeoutExpired when the operation timed out 77 | void timeoutTriggered() override; 78 | 79 | // Our operation has completed. During completeOperation, 80 | // specializedCompleteOperation is invoked for subclasses to perform 81 | // their own finalization (typically annotating errors and handling 82 | // timeouts). 83 | void completeOperation(OperationResult result); 84 | void completeOperationInner(OperationResult result); 85 | virtual void specializedTimeoutTriggered() = 0; 86 | virtual void specializedCompleteOperation() = 0; 87 | 88 | void protocolCompleteOperation(OperationResult result) override; 89 | 90 | bool isInEventBaseThread() const; 91 | bool isEventBaseSet() const; 92 | 93 | // This will contain the max block time of the thread 94 | Duration max_thread_block_time_ = Duration(0); 95 | Duration total_thread_block_time_ = Duration(0); 96 | 97 | // Friends because they need to access the query callbacks on this class 98 | template 99 | friend folly::SemiFuture handlePreQueryCallback(Operation& op); 100 | template 101 | friend void handleQueryCompletion( 102 | Operation& op, 103 | QueryResult query_result, 104 | QueryCallbackReason reason, 105 | folly::Promise>& promise); 106 | 107 | private: 108 | // Restore folly::RequestContext and also invoke actionable() 109 | void invokeActionable(); 110 | 111 | friend class Operation; 112 | friend class Connection; 113 | friend class ConnectOperation; 114 | friend class SyncConnection; 115 | friend class SyncConnectionPool; 116 | }; 117 | 118 | } // namespace facebook::common::mysql_client::mysql_protocol 119 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlResult.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/mysql_protocol/MysqlResult.h" 10 | #include "squangle/mysql_client/mysql_protocol/MysqlConnection.h" 11 | #include "squangle/mysql_client/mysql_protocol/MysqlRow.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlRowMetadata.h" 13 | 14 | namespace facebook::common::mysql_client::mysql_protocol { 15 | 16 | class MysqlRowMetadata; 17 | 18 | size_t MysqlResult::numFields() const { 19 | DCHECK(res_); 20 | auto ret = mysql_num_fields(res_.get()); 21 | VLOG(4) << fmt::format( 22 | "mysql_num_fields({}) returned {}", (void*)res_.get(), ret); 23 | return ret; 24 | } 25 | 26 | MYSQL_FIELD* MysqlResult::fields() const { 27 | DCHECK(res_); 28 | auto ret = mysql_fetch_fields(res_.get()); 29 | VLOG(4) << fmt::format( 30 | "mysql_fetch_fields({}) returned {}", (void*)res_.get(), (void*)ret); 31 | return ret; 32 | } 33 | 34 | size_t MysqlResult::numRows() const { 35 | DCHECK(res_); 36 | auto ret = mysql_num_rows(res_.get()); 37 | VLOG(4) << fmt::format( 38 | "mysql_num_rows({}) returned {}", (void*)res_.get(), ret); 39 | return ret; 40 | } 41 | 42 | std::unique_ptr MysqlResult::getRowMetadata() const { 43 | return std::make_unique(*this); 44 | } 45 | 46 | namespace { 47 | 48 | InternalResult::FetchRowRet MysqlRowFactory( 49 | MYSQL_RES* result, 50 | MYSQL_ROW mysqlRow) { 51 | std::unique_ptr row; 52 | if (mysqlRow) { 53 | auto* lengths = mysql_fetch_lengths(result); 54 | VLOG(4) << fmt::format( 55 | "mysql_fetch_lengths({}) returned {}", (void*)result, (void*)lengths); 56 | 57 | auto numFields = mysql_num_fields(result); 58 | VLOG(4) << fmt::format( 59 | "mysql_num_fields({}) returned {}", (void*)result, numFields); 60 | 61 | row = std::make_unique(mysqlRow, numFields, lengths); 62 | } 63 | 64 | return std::make_pair(DONE, std::move(row)); 65 | } 66 | 67 | } // namespace 68 | 69 | InternalResult::FetchRowRet SyncMysqlResult::fetchRow() { 70 | DCHECK(res_); 71 | auto mysqlRow = mysql_fetch_row(res_.get()); 72 | VLOG(4) << fmt::format( 73 | "mysql_fetch_row({}) returned {}", (void*)res_.get(), (void*)mysqlRow); 74 | 75 | return MysqlRowFactory(res_.get(), mysqlRow); 76 | } 77 | 78 | InternalResult::FetchRowRet AsyncMysqlResult::fetchRow() { 79 | std::unique_ptr row; 80 | DCHECK(res_); 81 | MYSQL_ROW mysqlRow; 82 | auto ret = mysql_fetch_row_nonblocking(res_.get(), &mysqlRow); 83 | VLOG(4) << fmt::format( 84 | "mysql_fetch_row_nonblocking({}) returned {}, MYSQL_ROW = {}", 85 | (void*)res_.get(), 86 | ret, 87 | (void*)mysqlRow); 88 | 89 | if (ret == NET_ASYNC_COMPLETE) { 90 | return MysqlRowFactory(res_.get(), mysqlRow); 91 | } 92 | 93 | return std::make_pair(MysqlConnection::toHandlerStatus(ret), std::move(row)); 94 | } 95 | 96 | } // namespace facebook::common::mysql_client::mysql_protocol 97 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlResult.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/InternalConnection.h" 12 | 13 | namespace facebook::common::mysql_client::mysql_protocol { 14 | 15 | class MysqlResult : public InternalResult { 16 | public: 17 | explicit MysqlResult(MYSQL_RES* res) : res_(res) {} 18 | 19 | [[nodiscard]] size_t numFields() const; 20 | 21 | [[nodiscard]] MYSQL_FIELD* fields() const; 22 | 23 | [[nodiscard]] size_t numRows() const override; 24 | 25 | void close() override { 26 | res_->handle = nullptr; 27 | } 28 | 29 | protected: 30 | friend class MysqlFetchOperationImpl; 31 | 32 | using MysqlResultDeleter = 33 | folly::static_function_deleter; 34 | using MysqlResultPtr = std::unique_ptr; 35 | 36 | [[nodiscard]] std::unique_ptr getRowMetadata() const; 37 | 38 | MysqlResultPtr res_; 39 | }; 40 | 41 | class AsyncMysqlResult : public MysqlResult { 42 | public: 43 | explicit AsyncMysqlResult(MYSQL_RES* res) : MysqlResult(res) {} 44 | 45 | [[nodiscard]] FetchRowRet fetchRow() override; 46 | }; 47 | 48 | class SyncMysqlResult : public MysqlResult { 49 | public: 50 | explicit SyncMysqlResult(MYSQL_RES* res) : MysqlResult(res) {} 51 | 52 | [[nodiscard]] FetchRowRet fetchRow() override; 53 | }; 54 | 55 | } // namespace facebook::common::mysql_client::mysql_protocol 56 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlRow.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/InternalConnection.h" 12 | 13 | namespace facebook::common::mysql_client::mysql_protocol { 14 | 15 | class MysqlRow : public InternalRow { 16 | public: 17 | MysqlRow(MYSQL_ROW row, size_t numCols, unsigned long* lengths) 18 | : row_(std::move(row)), numCols_(numCols), lengths_(lengths) { 19 | DCHECK(row_); 20 | DCHECK(numCols_); 21 | DCHECK(lengths_); 22 | } 23 | 24 | [[nodiscard]] folly::StringPiece columnString(size_t col) const override { 25 | DCHECK_LT(col, numCols_); 26 | DCHECK(row_[col]); 27 | return folly::StringPiece(row_[col], lengths_[col]); 28 | } 29 | 30 | [[nodiscard]] InternalRow::Type columnType(size_t col) const override { 31 | DCHECK_LT(col, numCols_); 32 | if (!row_[col]) { 33 | return InternalRow::Type::Null; 34 | } 35 | 36 | // MySQL returns all data as strings - note this is NOT the column type - 37 | // just the type of the data returned via this protocol 38 | return InternalRow::Type::String; 39 | } 40 | 41 | [[nodiscard]] size_t columnLength(size_t col) const override { 42 | DCHECK_LT(col, numCols_); 43 | return lengths_[col]; 44 | } 45 | 46 | private: 47 | MYSQL_ROW row_; 48 | size_t numCols_; 49 | unsigned long* lengths_; 50 | }; 51 | 52 | } // namespace facebook::common::mysql_client::mysql_protocol 53 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlRowMetadata.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/InternalConnection.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlResult.h" 13 | 14 | namespace facebook::common::mysql_client::mysql_protocol { 15 | 16 | class MysqlRowMetadata : public InternalRowMetadata { 17 | public: 18 | explicit MysqlRowMetadata(const MysqlResult& result) 19 | : num_fields_(result.numFields()), fields_(result.fields()) {} 20 | 21 | [[nodiscard]] size_t numFields() const noexcept override { 22 | return num_fields_; 23 | } 24 | 25 | [[nodiscard]] folly::StringPiece getTableName(size_t index) const override { 26 | DCHECK_LT(index, num_fields_); 27 | return folly::StringPiece( 28 | fields_[index].table, fields_[index].table_length); 29 | } 30 | 31 | [[nodiscard]] folly::StringPiece getFieldName(size_t index) const override { 32 | DCHECK_LT(index, num_fields_); 33 | return folly::StringPiece(fields_[index].name, fields_[index].name_length); 34 | } 35 | 36 | [[nodiscard]] enum_field_types getFieldType(size_t index) const override { 37 | DCHECK_LT(index, num_fields_); 38 | return fields_[index].type; 39 | } 40 | 41 | [[nodiscard]] uint64_t getFieldFlags(size_t index) const override { 42 | DCHECK_LT(index, num_fields_); 43 | return fields_[index].flags; 44 | } 45 | 46 | private: 47 | size_t num_fields_; 48 | const MYSQL_FIELD* fields_; 49 | }; 50 | 51 | } // namespace facebook::common::mysql_client::mysql_protocol 52 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlSpecialOperationImpl.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #include "squangle/mysql_client/mysql_protocol/MysqlSpecialOperationImpl.h" 10 | #include "squangle/mysql_client/Connection.h" 11 | #include "squangle/mysql_client/mysql_protocol/MysqlConnection.h" 12 | 13 | namespace facebook::common::mysql_client::mysql_protocol { 14 | 15 | void MysqlSpecialOperationImpl::actionable() { 16 | auto status = runSpecialOperation(); 17 | if (status == PENDING) { 18 | waitForActionable(); 19 | } else { 20 | auto result = (status == DONE) ? OperationResult::Succeeded 21 | : OperationResult::Failed; // ERROR 22 | completeOperation(result); 23 | if (callback_) { 24 | callback_(getOp(), result); 25 | } 26 | } 27 | } 28 | 29 | void MysqlSpecialOperationImpl::specializedCompleteOperation() { 30 | conn().notify(); 31 | } 32 | 33 | void MysqlSpecialOperationImpl::specializedTimeoutTriggered() { 34 | completeOperation(OperationResult::TimedOut); 35 | } 36 | 37 | void MysqlSpecialOperationImpl::specializedRun() { 38 | const auto* mysql_conn = getMysqlConnection(); 39 | changeHandlerFD( 40 | folly::NetworkSocket::fromFd(mysql_conn->getSocketDescriptor())); 41 | actionable(); 42 | } 43 | 44 | } // namespace facebook::common::mysql_client::mysql_protocol 45 | -------------------------------------------------------------------------------- /squangle/mysql_client/mysql_protocol/MysqlSpecialOperationImpl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include "squangle/mysql_client/SpecialOperation.h" 12 | #include "squangle/mysql_client/mysql_protocol/MysqlOperationImpl.h" 13 | 14 | namespace facebook::common::mysql_client::mysql_protocol { 15 | 16 | class MysqlSpecialOperationImpl : public MysqlOperationImpl, 17 | public SpecialOperationImpl { 18 | public: 19 | explicit MysqlSpecialOperationImpl(std::unique_ptr conn) 20 | : OperationBase(std::move(conn)) {} 21 | 22 | protected: 23 | void actionable() override; 24 | void specializedCompleteOperation() override; 25 | void specializedTimeoutTriggered() override; 26 | void specializedRun() override; 27 | }; 28 | 29 | } // namespace facebook::common::mysql_client::mysql_protocol 30 | -------------------------------------------------------------------------------- /squangle/util/StorageRow.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace facebook::common::mysql_client { 17 | 18 | /** 19 | The StorageRow class is an attempt to store MySQL row data in a space efficient 20 | manner that works for both the MySQL protocol (data always returned as strings) 21 | and the Thrift protocol (data returned in native format). 22 | 23 | Each entry has at minimum an entry in `offset_`. The high bit of the offset is 24 | set if the column is null. Otherwise the other 31 bits are the offset into 25 | `data_` where the data for the column starts. 26 | 27 | Assuming the data is not null, information about the column's contents is 28 | then written into `data_`. The first byte is a combination of 2 bits for the 29 | basic type (bool, integer, double, string) and 6 bits for flags. For bool data, 30 | the value is stored as a flag in the control flags and no other data is 31 | necessary. 32 | 33 | For integers we store whether the data was written as signed or unsigned as a 34 | flag in the control flags. Two more bits are used to store whether the data is 35 | written as a single byte, two bytes, four bytes or eight bytes. We could 36 | support more versions (3 bytes, 5 bytes, etc.) but the code to handle these is 37 | more complicated so I stuck to the standard sizes. 38 | 39 | Doubles are stored natively. There isn't any compression we can do that isn't 40 | lossy. 41 | 42 | Strings are stored one of two ways with a flag in the control flags used to 43 | differentiate the two. Strings of 4096 bytes or less are stored straight in 44 | `data_`, byte for byte (without the null terminator). Strings longer than that 45 | have an index written into `data_` (using the variable integer format) and then 46 | are put into `long_strings_` (at the specified index) individually. The benefit 47 | of doing this is that we can make the type of `control_` a vector of `uint32_t` 48 | instead of `size_t` saving 4 bytes per column. 49 | **/ 50 | 51 | class StorageRow { 52 | private: 53 | static constexpr auto kTypeBits = 3; 54 | static constexpr auto kOffsetBits = 32 - kTypeBits; 55 | static constexpr auto kOffsetMask = (1U << kOffsetBits) - 1; 56 | static constexpr auto kSmallStringLimit = 4096; 57 | 58 | // Used for variable sized integers (used for 'long' strings) 59 | static constexpr auto kVarUIntContinuation = (std::byte)0x80; 60 | static constexpr auto kVarUIntMask = (std::byte)0x7F; 61 | static constexpr auto kVarUIntShift = 7; 62 | 63 | public: 64 | explicit StorageRow(size_t expectedSize); 65 | 66 | void appendValue(folly::StringPiece data); 67 | void appendValue(const char* str) { 68 | appendValue(folly::StringPiece(str)); 69 | } 70 | void appendValue(bool data); 71 | void appendValue(uint64_t data); 72 | void appendValue(int64_t data); 73 | void appendValue(double data); 74 | 75 | void appendNull(); 76 | 77 | size_t count() const { 78 | return control_.size(); 79 | } 80 | 81 | bool isNull(size_t column) const; 82 | 83 | template 84 | T as(size_t column, Func&& func) const { 85 | DCHECK_LT(column, control_.size()); 86 | 87 | Offsets offsets; 88 | auto data_type = readControl(column, &offsets); 89 | // Double-check that all the bytes were consumed when we are done. If not 90 | // that means we have a bug in the encoding protocol. 91 | auto guard = 92 | folly::makeGuard([&] { DCHECK_EQ(offsets.curr, offsets.end); }); 93 | 94 | switch (data_type) { 95 | case dtNull: 96 | return func(folly::StringPiece()); 97 | 98 | case dtBoolTrue: 99 | return func(true); 100 | 101 | case dtBoolFalse: 102 | return func(false); 103 | 104 | case dtUnsignedInteger: 105 | return func(readUnsignedInteger(offsets)); 106 | 107 | case dtSignedInteger: 108 | return func(readSignedInteger(offsets)); 109 | 110 | case dtDouble: 111 | return func(readDouble(offsets)); 112 | 113 | case dtShortString: 114 | return func(readShortString(offsets)); 115 | 116 | case dtLongString: 117 | return func(readLongString(offsets)); 118 | 119 | default: 120 | DCHECK(false); 121 | throw std::runtime_error("Invalid data type"); 122 | } 123 | 124 | folly::assume_unreachable(); 125 | } 126 | 127 | private: 128 | // With only 8 values this enum can fit in 3 bits. The `control_` member is a 129 | // vector of `uint32_t` where the high 3 bits represent values in this enum 130 | // and the low 29 bits represent the offset into `data_` for the column. 131 | enum DataType { 132 | dtNull, 133 | dtBoolTrue, 134 | dtBoolFalse, 135 | dtSignedInteger, 136 | dtUnsignedInteger, 137 | dtDouble, 138 | dtShortString, 139 | dtLongString, 140 | }; 141 | 142 | struct Offsets { 143 | uint32_t curr; 144 | uint32_t end; 145 | }; 146 | 147 | void writeControl(DataType data_type); 148 | 149 | void writeByte(std::byte byte); 150 | void writeDouble(double value); 151 | void writeBytes(const void* data, size_t size); 152 | void writeBytes(folly::StringPiece str); 153 | void writeVarUInt(uint64_t var); 154 | 155 | DataType readControl(size_t column, Offsets* offsets) const; 156 | std::byte readByte(Offsets& offsets) const; 157 | uint64_t readUnsignedInteger(Offsets& offsets) const; 158 | int64_t readSignedInteger(Offsets& offsets) const; 159 | uint64_t readVarUInt(Offsets& offsets) const; 160 | 161 | template 162 | T readBytes(Offsets& offsets) const { 163 | DCHECK_LE(offsets.curr + sizeof(T), offsets.end); 164 | T res; 165 | memcpy(&res, &data_[offsets.curr], sizeof(res)); 166 | offsets.curr += sizeof(res); 167 | return res; 168 | } 169 | 170 | double readDouble(Offsets& offsets) const; 171 | folly::StringPiece readShortString(Offsets& offsets) const; 172 | folly::StringPiece readLongString(Offsets& offsets) const; 173 | 174 | // This holds the offsets into `data_` for each column in this row. Every 175 | // column has to be represented (even if null). We can use a uint32_t here 176 | // because we pull large strings into a separate location. 31 bits (1 bit is 177 | // used for null values) gives us 2Gb of space per row. Strings longer than 178 | // 4Kb stored elsewhere, so in the worst case we would store 4K (column_size) 179 | // * 4K (maximum_cols_per_table) for about 16Mb in data_. 180 | std::vector control_; 181 | // This holds all the basic data, not including long strings) for the row. The 182 | // data for any particular column is accessed via `control_[column]` and the 183 | // number of bytes for that column can be derived by subtracting the offset 184 | // from the next column's offset. 185 | std::vector data_; 186 | 187 | // This holds any long (>4096 bytes) strings. It is more efficient to hold 188 | // them outside of `data_` as this allows us to use 32 bits for offsets per 189 | // column instead of 64. Use a unique pointer for this vector so that when we 190 | // don't need any long strings we only take up 8 bytes for the unique_ptr vs. 191 | // 24 for a vector. 192 | using LongStringsVector = std::vector; 193 | std::unique_ptr long_strings_; 194 | }; 195 | 196 | } // namespace facebook::common::mysql_client 197 | -------------------------------------------------------------------------------- /squangle/util/StringStore.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * All rights reserved. 4 | * 5 | * This source code is licensed under the BSD-style license found in the 6 | * LICENSE file in the root directory of this source tree. 7 | */ 8 | 9 | #pragma once 10 | 11 | #include 12 | #include 13 | 14 | namespace facebook::common::mysql_client { 15 | 16 | template 17 | class StringStore { 18 | public: 19 | using ValueGenerator = std::function; 20 | const std::string& getString(Key&& key, const ValueGenerator& gen) { 21 | // See if we already have a string for the specified key 22 | if (auto it = map_.find(key); it != map_.end()) { 23 | // If so, return it. 24 | return it->second; 25 | } 26 | 27 | // If not, create a new string and insert it into the map. 28 | auto [it, _] = map_.emplace(std::move(key), gen()); 29 | return it->second; 30 | } 31 | 32 | private: 33 | // Whatever map we use for this must guarantee reference stability as 34 | // `getString` returns a reference to values in the map. 35 | std::unordered_map map_; 36 | }; 37 | 38 | } // namespace facebook::common::mysql_client 39 | --------------------------------------------------------------------------------