Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add client max connection #136

Open
wants to merge 18 commits into
base: unstable
Choose a base branch
from
Open

feat: add client max connection #136

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
eaa8fe8
feat: addClientMaxConnectionLimits(#104)
Ph0m1 Dec 10, 2024
e01a031
feat: addClientsMaxConnectionLimits
Ph0m1 Dec 11, 2024
a3d3ba3
remove unused header
Ph0m1 Dec 21, 2024
ab5c2c2
feat: rewrite base on netoption
Ph0m1 Dec 22, 2024
3e07a20
add response to client
Ph0m1 Jan 10, 2025
9edb5b4
Merge branch 'unstable' of github.com:arana-db/kiwi into feat/add_cli…
Ph0m1 Jan 10, 2025
80da4ab
update format
Ph0m1 Jan 10, 2025
e555f07
after clang format
Ph0m1 Jan 10, 2025
1743986
after clang format
Ph0m1 Jan 10, 2025
e79f013
after clang format
Ph0m1 Jan 10, 2025
81a60af
Merge branch 'unstable' of github.com:arana-db/kiwi into feat/add_cli…
Ph0m1 Jan 18, 2025
92fec59
merge with unstable
Ph0m1 Jan 18, 2025
30a2327
rewrite codes with workflows
Ph0m1 Jan 18, 2025
c7b7c79
after clang-format
Ph0m1 Jan 18, 2025
239f72f
after clang-format
Ph0m1 Jan 18, 2025
50a3929
after clang-format
Ph0m1 Jan 18, 2025
4e9c500
remove increment client count function
Ph0m1 Jan 18, 2025
3c5b37f
fix format od INFO
Ph0m1 Jan 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/kiwi.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ bool KiwiDB::Init() {
auto num = g_config.worker_threads_num + g_config.slave_threads_num;
options_.SetThreadNum(num);

options_.SetMaxClients(g_config.max_clients);

// now we only use fast cmd thread pool
auto status = cmd_threads_.Init(g_config.fast_cmd_threads_num, 1, "kiwi-cmd");
if (!status.ok()) {
Expand Down
1 change: 1 addition & 0 deletions src/net/event_server.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ class EventServer final {
std::condition_variable cv_;

std::shared_ptr<Timer> timer_;

};

template <typename T>
Expand Down
6 changes: 6 additions & 0 deletions src/net/net_options.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,16 @@ class NetOptions {

bool GetRwSeparation() const { return rwSeparation_; }

void SetMaxClients(uint32_t maxClients) { maxClients_ = maxClients; }

uint32_t GetMaxClients() const { return maxClients_; }

private:
bool rwSeparation_ = true; // Whether to separate read and write

int8_t threadNum_ = 1; // The number of threads

uint32_t maxClients_ = 1; // The maximum number of connections(default 40000)
};

} // namespace net
35 changes: 29 additions & 6 deletions src/net/thread_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,17 +90,25 @@ class ThreadManager {

uint64_t DoTCPConnect(T &t, int fd, const std::shared_ptr<Connection> &conn);

uint32_t getClientCount() const { return clientCount_.load(); }

void clientCountIncrement() { clientCount_.fetch_add(1, std::memory_order_relaxed); }

void clientCountDecrement() { clientCount_.fetch_sub(1, std::memory_order_relaxed); }

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Maintain consistent naming convention and review memory ordering.

  1. Method names should follow the class's snake_case convention:

    • getClientCount → get_client_count
    • clientCountIncrement → client_count_increment
    • clientCountDecrement → client_count_decrement

  2. Consider using std::memory_order_seq_cst instead of memory_order_relaxed for the atomic operations since the client count is used for critical decision-making in OnNetEventCreate.

-  uint32_t getClientCount() const { return clientCount_.load(); }
+  uint32_t get_client_count() const { return client_count_.load(); }

-  void clientCountIncrement() { clientCount_.fetch_add(1, std::memory_order_relaxed); }
+  void client_count_increment() { client_count_.fetch_add(1, std::memory_order_seq_cst); }

-  void clientCountDecrement() { clientCount_.fetch_sub(1, std::memory_order_relaxed); }
+  void client_count_decrement() { client_count_.fetch_sub(1, std::memory_order_seq_cst); }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
uint32_t getClientCount() const { return clientCount_.load(); }
void clientCountIncrement() { clientCount_.fetch_add(1, std::memory_order_relaxed); }
void clientCountDecrement() { clientCount_.fetch_sub(1, std::memory_order_relaxed); }
uint32_t get_client_count() const { return client_count_.load(); }
void client_count_increment() { client_count_.fetch_add(1, std::memory_order_seq_cst); }
void client_count_decrement() { client_count_.fetch_sub(1, std::memory_order_seq_cst); }
🧰 Tools
🪛 GitHub Actions: kiwi

[error] File has clang-format style issues

private:
const int8_t index_ = 0; // The index of the thread
std::atomic<bool> running_ = true; // Whether the thread is running

NetOptions netOptions_;

inline static std::atomic<uint32_t> clientCount_{0};
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codebase verification

Fix inconsistent memory ordering in client count operations

The client count operations use inconsistent memory ordering:

  • compare_exchange_strong uses memory_order_seq_cst
  • fetch_sub uses memory_order_relaxed

This asymmetry could lead to temporary inconsistencies in client counting. Consider using consistent ordering:

void client_count_decrement() { 
    clientCount_.fetch_sub(1, std::memory_order_seq_cst);
}
🔗 Analysis chain

Consider memory ordering implications of static initialization.

The static atomic counter is correctly used for thread-safe counting, but consider potential initialization order issues across translation units. Also, the relaxed memory ordering in client_count_decrement might be too weak for accurate client counting.

Run this script to check for other static atomics that might interact with this counter:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for other static atomic variables that might have initialization order dependencies
rg -l 'static\s+std::atomic'

# Search for other client count related variables
ast-grep --pattern 'std::atomic<$_> $_ = $_'

Length of output: 130

Script:

#!/bin/bash
# Check content of log_index.h for static atomics
echo "=== log_index.h static atomics ==="
rg -A 2 -B 2 'static\s+std::atomic' src/storage/src/log_index.h

# Check memory ordering usage
echo -e "\n=== Memory ordering usage ==="
rg 'memory_order|\.load\(|\.store\(|\.fetch_' src/net/thread_manager.h src/storage/src/log_index.h

# Check for any direct interactions between these files
echo -e "\n=== File interactions ==="
rg -l "thread_manager.*log_index|log_index.*thread_manager"

Length of output: 2006


std::unique_ptr<IOThread> readThread_; // Read thread
std::unique_ptr<IOThread> writeThread_; // Write thread

// All connections for the current thread
std::unordered_map<uint64_t, std::pair<T, std::shared_ptr<Connection>>> connections_;
std::unordered_map<uint64_t, std::pair<T, std::shared_ptr<Connection>>>
connections_; // All connections for the current thread

std::shared_mutex mutex_;

Expand All @@ -116,7 +124,10 @@ class ThreadManager {
};

template <typename T>
requires HasSetFdFunction<T> ThreadManager<T>::~ThreadManager() { Stop(); }
requires HasSetFdFunction<T>
ThreadManager<T>::~ThreadManager() {
Stop();
}

template <typename T>
requires HasSetFdFunction<T>
Expand Down Expand Up @@ -145,6 +156,14 @@ void ThreadManager<T>::Stop() {
template <typename T>
requires HasSetFdFunction<T>
void ThreadManager<T>::OnNetEventCreate(int fd, const std::shared_ptr<Connection> &conn) {
if (getClientCount() >= netOptions_.GetMaxClients()) {
INFO("Max client connetions, refuse new connection fd: %d", fd);
std::string response = "-ERR max clients reached\r\n";
::send(fd, response.c_str(), response.size(), 0);
::close(fd);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

添加日志

return;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

返回 redis 格式的一个 response

}

T t;
onInit_(&t);
auto connId = getConnId();
Expand All @@ -166,6 +185,7 @@ void ThreadManager<T>::OnNetEventCreate(int fd, const std::shared_ptr<Connection
}

onCreate_(connId, t, conn->addr_);
clientCountIncrement();
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix potential race condition in client counting.

There's a race condition between checking the client count and incrementing it. Another thread could add a client between the check and increment, potentially exceeding the max limit.

 template <typename T>
 requires HasSetFdFunction<T>
 void ThreadManager<T>::OnNetEventCreate(int fd, const std::shared_ptr<Connection> &conn) {
-  if (getClientCount() >= netOptions_.GetMaxClients()) {
+  if (!client_count_.compare_exchange_strong(
+          expected,
+          expected + 1,
+          std::memory_order_seq_cst,
+          std::memory_order_seq_cst) ||
+      expected >= netOptions_.GetMaxClients()) {
     INFO("Max client connections, refuse new connection fd: %d", fd);
     std::string response = "-ERR max clients reached\r\n";
     ::send(fd, response.c_str(), response.size(), 0);
     ::close(fd);
     return;
   }
   
   // ... rest of the connection setup ...
-  clientCountIncrement();  // Remove this as we've already incremented atomically
 }

Also applies to: 225-225

🧰 Tools
🪛 GitHub Actions: kiwi

[error] File has clang-format style issues

}

template <typename T>
Expand Down Expand Up @@ -202,11 +222,14 @@ void ThreadManager<T>::OnNetEventClose(uint64_t connId, std::string &&err) {
iter->second.second->netEvent_->Close(); // close socket
onClose_(iter->second.first, std::move(err));
connections_.erase(iter);
clientCountDecrement();
}

template <typename T>
requires HasSetFdFunction<T>
void ThreadManager<T>::CloseConnection(uint64_t connId) { OnNetEventClose(connId, ""); }
void ThreadManager<T>::CloseConnection(uint64_t connId) {
OnNetEventClose(connId, "");
}

template <typename T>
requires HasSetFdFunction<T>
Expand Down Expand Up @@ -330,8 +353,8 @@ bool ThreadManager<T>::CreateWriteThread() {
}

template <typename T>
requires HasSetFdFunction<T> uint64_t ThreadManager<T>::DoTCPConnect(T &t, int fd,
const std::shared_ptr<Connection> &conn) {
requires HasSetFdFunction<T>
uint64_t ThreadManager<T>::DoTCPConnect(T &t, int fd, const std::shared_ptr<Connection> &conn) {
auto connId = getConnId();
if constexpr (IsPointer_v<T>) {
t->SetConnId(connId);
Expand Down
Loading