519 lines
18 KiB
C
Raw Normal View History

Initial commit
2025-09-27 00:33:35 +02:00
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <errno.h>
#include <math.h>
#include <stdint.h>
// --- Test Configuration ---
#define HOST "127.0.0.1"
#define PORT 8080
#define NUM_SUBSCRIBERS 1000
#define NUM_PUBLISHERS 10
#define TOTAL_CLIENTS (NUM_SUBSCRIBERS + NUM_PUBLISHERS)
#define TEST_DURATION_S 15
#define MESSAGES_PER_SECOND_PER_PUBLISHER 100
// --- Internal Configuration ---
#define MAX_EVENTS TOTAL_CLIENTS
#define RW_BUFFER_SIZE 8192
#define MAX_LATENCIES 20000000 // Pre-allocate for ~1.3M messages/sec
// --- WebSocket Constants ---
#define WEBSOCKET_KEY_MAGIC "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
// --- Helper Enums ---
typedef enum {
CLIENT_SUBSCRIBER,
CLIENT_PUBLISHER
} ClientType;
typedef enum {
STATE_CONNECTING,
STATE_HANDSHAKE_SEND,
STATE_HANDSHAKE_RECV,
STATE_SUBSCRIBING,
STATE_RUNNING,
STATE_CLOSED
} ClientState;
// --- Client State Structure ---
typedef struct {
int fd;
ClientType type;
ClientState state;
char read_buf[RW_BUFFER_SIZE];
size_t read_len;
char write_buf[RW_BUFFER_SIZE];
size_t write_len;
size_t write_pos;
double next_send_time;
} Client;
// --- Global State & Metrics ---
double* latencies;
size_t latencies_count = 0;
uint64_t messages_sent = 0;
uint64_t messages_received = 0;
int subscriber_setup_count = 0;
int all_subscribed = 0;
int epoll_fd;
const char* CHANNELS[] = {"news", "sports", "tech", "finance", "weather"};
const int NUM_CHANNELS = sizeof(CHANNELS) / sizeof(CHANNELS[0]);
// =============================================================================
// START: Clean SHA-1 and Base64 Implementations
// =============================================================================
// --- SHA-1 Implementation ---
typedef struct {
uint32_t state[5];
uint32_t count[2];
unsigned char buffer[64];
} SHA1_CTX;
#define SHA1_ROTLEFT(n,c) (((n) << (c)) | ((n) >> (32 - (c))))
void SHA1_Transform(uint32_t state[5], const unsigned char buffer[64]) {
uint32_t a, b, c, d, e;
uint32_t block[16];
memcpy(block, buffer, 64);
for(int i = 0; i < 16; i++) {
uint8_t *p = (uint8_t*)&block[i];
block[i] = (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
}
a = state[0]; b = state[1]; c = state[2]; d = state[3]; e = state[4];
uint32_t W[80];
for (int t = 0; t < 80; t++) {
if (t < 16) {
W[t] = block[t];
} else {
W[t] = SHA1_ROTLEFT(W[t - 3] ^ W[t - 8] ^ W[t - 14] ^ W[t - 16], 1);
}
uint32_t temp = SHA1_ROTLEFT(a, 5) + e + W[t];
if (t < 20) temp += ((b & c) | (~b & d)) + 0x5A827999;
else if (t < 40) temp += (b ^ c ^ d) + 0x6ED9EBA1;
else if (t < 60) temp += ((b & c) | (b & d) | (c & d)) + 0x8F1BBCDC;
else temp += (b ^ c ^ d) + 0xCA62C1D6;
e = d; d = c; c = SHA1_ROTLEFT(b, 30); b = a; a = temp;
}
state[0] += a; state[1] += b; state[2] += c; state[3] += d; state[4] += e;
}
void SHA1_Init(SHA1_CTX* context) {
context->state[0] = 0x67452301;
context->state[1] = 0xEFCDAB89;
context->state[2] = 0x98BADCFE;
context->state[3] = 0x10325476;
context->state[4] = 0xC3D2E1F0;
context->count[0] = context->count[1] = 0;
}
void SHA1_Update(SHA1_CTX* context, const unsigned char* data, uint32_t len) {
uint32_t i, j;
j = context->count[0];
if ((context->count[0] += len << 3) < j) context->count[1]++;
context->count[1] += (len >> 29);
j = (j >> 3) & 63;
if ((j + len) > 63) {
memcpy(&context->buffer[j], data, (i = 64 - j));
SHA1_Transform(context->state, context->buffer);
for (; i + 63 < len; i += 64) {
SHA1_Transform(context->state, &data[i]);
}
j = 0;
} else {
i = 0;
}
memcpy(&context->buffer[j], &data[i], len - i);
}
void SHA1_Final(unsigned char digest[20], SHA1_CTX* context) {
uint32_t i;
unsigned char finalcount[8];
for (i = 0; i < 8; i++) {
finalcount[i] = (unsigned char)((context->count[(i >= 4 ? 0 : 1)] >> ((3 - (i & 3)) * 8)) & 255);
}
SHA1_Update(context, (unsigned char*)"\x80", 1);
while ((context->count[0] & 504) != 448) {
SHA1_Update(context, (unsigned char*)"\0", 1);
}
SHA1_Update(context, finalcount, 8);
for (i = 0; i < 20; i++) {
digest[i] = (unsigned char)((context->state[i >> 2] >> ((3 - (i & 3)) * 8)) & 255);
}
}
// --- Base64 Implementation ---
char* base64_encode(const unsigned char *data, size_t input_length) {
const char b64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
size_t output_length = 4 * ((input_length + 2) / 3);
char *encoded_data = malloc(output_length + 1);
if (encoded_data == NULL) return NULL;
for (size_t i = 0, j = 0; i < input_length;) {
uint32_t octet_a = i < input_length ? data[i++] : 0;
uint32_t octet_b = i < input_length ? data[i++] : 0;
uint32_t octet_c = i < input_length ? data[i++] : 0;
uint32_t triple = (octet_a << 16) + (octet_b << 8) + octet_c;
encoded_data[j++] = b64_table[(triple >> 18) & 0x3F];
encoded_data[j++] = b64_table[(triple >> 12) & 0x3F];
encoded_data[j++] = b64_table[(triple >> 6) & 0x3F];
encoded_data[j++] = b64_table[triple & 0x3F];
}
for (size_t i = 0; i < (3 - input_length % 3) % 3; i++) {
encoded_data[output_length - 1 - i] = '=';
}
encoded_data[output_length] = '\0';
return encoded_data;
}
// =============================================================================
// END: Clean SHA-1 and Base64 Implementations
// =============================================================================
// --- Utility Functions ---
double get_time_double() {
struct timespec ts;
clock_gettime(CLOCK_REALTIME, &ts);
return ts.tv_sec + ts.tv_nsec / 1e9;
}
void epoll_ctl_mod(int fd, uint32_t events, void* ptr) {
struct epoll_event ev;
ev.events = events;
ev.data.ptr = ptr;
epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &ev);
}
void close_client(Client* client) {
if (client->state != STATE_CLOSED) {
epoll_ctl(epoll_fd, EPOLL_CTL_DEL, client->fd, NULL);
close(client->fd);
client->state = STATE_CLOSED;
}
}
// --- WebSocket Core Functions ---
size_t create_ws_frame(const char* payload, size_t payload_len, char* out_buffer) {
size_t frame_len = 2 + payload_len + 4; // Header + Mask + Payload
if (payload_len > 125) frame_len += 2; // For 16-bit length
out_buffer[0] = 0x81; // FIN + Text Frame
if (payload_len <= 125) {
out_buffer[1] = 0x80 | payload_len;
} else {
out_buffer[1] = 0x80 | 126;
*(uint16_t*)(out_buffer + 2) = htons(payload_len);
}
size_t header_len = (payload_len <= 125) ? 2 : 4;
uint32_t mask = rand();
*(uint32_t*)(out_buffer + header_len) = mask;
uint8_t* mask_bytes = (uint8_t*)&mask;
for(size_t i = 0; i < payload_len; ++i) {
out_buffer[header_len + 4 + i] = payload[i] ^ mask_bytes[i % 4];
}
return header_len + 4 + payload_len;
}
void send_handshake(Client* client) {
unsigned char key_bytes[16];
for (int i = 0; i < 16; i++) key_bytes[i] = rand() % 256;
char* b64_key = base64_encode(key_bytes, 16);
client->write_len = snprintf(client->write_buf, RW_BUFFER_SIZE,
"GET / HTTP/1.1\r\n"
"Host: %s:%d\r\n"
"Upgrade: websocket\r\n"
"Connection: Upgrade\r\n"
"Sec-WebSocket-Key: %s\r\n"
"Sec-WebSocket-Version: 13\r\n\r\n",
HOST, PORT, b64_key);
client->state = STATE_HANDSHAKE_SEND;
epoll_ctl_mod(client->fd, EPOLLIN | EPOLLOUT | EPOLLET, client);
free(b64_key);
}
// --- Event Handlers ---
void handle_write(Client* client) {
ssize_t sent = send(client->fd, client->write_buf + client->write_pos, client->write_len - client->write_pos, 0);
if (sent < 0) {
if (errno != EAGAIN && errno != EWOULDBLOCK) close_client(client);
return;
}
client->write_pos += sent;
if (client->write_pos >= client->write_len) {
client->write_pos = 0;
client->write_len = 0;
epoll_ctl_mod(client->fd, EPOLLIN | EPOLLET, client); // Done writing, wait for reads
if (client->state == STATE_HANDSHAKE_SEND) {
client->state = STATE_HANDSHAKE_RECV;
} else if (client->state == STATE_SUBSCRIBING) {
client->state = STATE_RUNNING;
subscriber_setup_count++;
if (subscriber_setup_count == NUM_SUBSCRIBERS) {
printf("✅ All subscribers are connected and subscribed. Starting publishers...\n");
all_subscribed = 1;
}
}
}
}
void handle_read(Client* client) {
ssize_t n = read(client->fd, client->read_buf + client->read_len, RW_BUFFER_SIZE - client->read_len);
if (n <= 0) {
if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) perror("read error");
close_client(client);
return;
}
client->read_len += n;
if (client->state == STATE_HANDSHAKE_RECV) {
if (strstr(client->read_buf, "\r\n\r\n")) {
if (strstr(client->read_buf, " 101 ") == NULL) {
fprintf(stderr, "Handshake failed for fd %d\n", client->fd);
close_client(client);
return;
}
client->read_len = 0; // Handshake complete, clear buffer
if (client->type == CLIENT_SUBSCRIBER) {
const char* channel = CHANNELS[rand() % NUM_CHANNELS];
char sub_msg[128];
int msg_len = snprintf(sub_msg, sizeof(sub_msg), "sub %s", channel);
client->write_len = create_ws_frame(sub_msg, msg_len, client->write_buf);
client->state = STATE_SUBSCRIBING;
epoll_ctl_mod(client->fd, EPOLLIN | EPOLLOUT | EPOLLET, client);
} else { // Publisher
client->state = STATE_RUNNING;
}
}
} else if (client->state == STATE_RUNNING && client->type == CLIENT_SUBSCRIBER) {
// Simple WebSocket frame parsing for this specific test case
while (client->read_len >= 2) {
uint64_t payload_len = client->read_buf[1] & 0x7F;
size_t header_len = 2;
if (payload_len == 126) {
if (client->read_len < 4) break;
payload_len = ntohs(*(uint16_t*)(client->read_buf + 2));
header_len = 4;
} else if (payload_len == 127) {
// Not expected for this test, would require 64-bit length handling
close_client(client);
break;
}
if (client->read_len >= header_len + payload_len) {
char* payload = client->read_buf + header_len;
double sent_time = atof(payload);
if (sent_time > 0) {
double latency = get_time_double() - sent_time;
if (latencies_count < MAX_LATENCIES) {
latencies[latencies_count++] = latency;
}
messages_received++;
}
size_t frame_size = header_len + payload_len;
memmove(client->read_buf, client->read_buf + frame_size, client->read_len - frame_size);
client->read_len -= frame_size;
} else {
break; // Incomplete frame
}
}
}
}
// --- Statistics ---
int compare_doubles(const void* a, const void* b) {
double da = *(const double*)a;
double db = *(const double*)b;
if (da < db) return -1;
if (da > db) return 1;
return 0;
}
void print_report() {
printf("\n"
"================================================================================\n");
printf("%s\n", " PERFORMANCE REPORT ");
printf("================================================================================\n");
if (latencies_count == 0) {
printf("No messages were received. Cannot generate a report. Is the server running?\n");
return;
}
uint64_t message_loss = (messages_sent > messages_received) ? (messages_sent - messages_received) : 0;
double loss_rate = (messages_sent > 0) ? ((double)message_loss / messages_sent * 100.0) : 0;
double throughput = (double)messages_received / TEST_DURATION_S;
printf("Test Duration: %d seconds\n", TEST_DURATION_S);
printf("Total Messages Sent: %lu\n", messages_sent);
printf("Total Messages Rcvd: %lu\n", messages_received);
printf("Message Loss: %lu (%.2f%%)\n", message_loss, loss_rate);
printf("Actual Throughput: %.2f msg/sec\n", throughput);
printf("--------------------------------------------------------------------------------\n");
qsort(latencies, latencies_count, sizeof(double), compare_doubles);
double sum = 0;
for (size_t i = 0; i < latencies_count; ++i) sum += latencies[i];
printf("Latency Statistics (ms):\n");
printf(" Average: %.4f ms\n", (sum / latencies_count) * 1000.0);
printf(" Min: %.4f ms\n", latencies[0] * 1000.0);
printf(" Max: %.4f ms\n", latencies[latencies_count - 1] * 1000.0);
printf(" Median (p50): %.4f ms\n", latencies[(size_t)(latencies_count * 0.50)] * 1000.0);
printf(" 95th Percentile: %.4f ms\n", latencies[(size_t)(latencies_count * 0.95)] * 1000.0);
printf(" 99th Percentile: %.4f ms\n", latencies[(size_t)(latencies_count * 0.99)] * 1000.0);
printf("================================================================================\n");
}
// --- Main Function ---
int main() {
srand(time(NULL));
latencies = malloc(sizeof(double) * MAX_LATENCIES);
if (!latencies) {
perror("malloc latencies");
return 1;
}
printf("Starting WebSocket Pub/Sub Load Test...\n");
printf("Simulating %d subscribers and %d publishers.\n", NUM_SUBSCRIBERS, NUM_PUBLISHERS);
printf("Publishing at ~%d msg/sec for %d seconds.\n", NUM_PUBLISHERS * MESSAGES_PER_SECOND_PER_PUBLISHER, TEST_DURATION_S);
printf("--------------------------------------------------------------------------------\n");
epoll_fd = epoll_create1(0);
if (epoll_fd == -1) {
perror("epoll_create1");
free(latencies);
return 1;
}
// *** FIX: Allocate clients on the heap, not the stack ***
Client* clients = malloc(sizeof(Client) * TOTAL_CLIENTS);
if (!clients) {
perror("malloc clients");
free(latencies);
close(epoll_fd);
return 1;
}
struct sockaddr_in server_addr;
server_addr.sin_family = AF_INET;
server_addr.sin_port = htons(PORT);
inet_pton(AF_INET, HOST, &server_addr.sin_addr);
for (int i = 0; i < TOTAL_CLIENTS; ++i) {
clients[i].fd = socket(AF_INET, SOCK_STREAM, 0);
fcntl(clients[i].fd, F_SETFL, O_NONBLOCK);
clients[i].type = (i < NUM_SUBSCRIBERS) ? CLIENT_SUBSCRIBER : CLIENT_PUBLISHER;
clients[i].state = STATE_CONNECTING;
clients[i].read_len = 0;
clients[i].write_len = 0;
clients[i].write_pos = 0;
clients[i].next_send_time = 0;
connect(clients[i].fd, (struct sockaddr*)&server_addr, sizeof(server_addr));
struct epoll_event ev;
ev.events = EPOLLIN | EPOLLOUT | EPOLLET;
ev.data.ptr = &clients[i];
if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, clients[i].fd, &ev) == -1) {
perror("epoll_ctl: add");
free(latencies);
free(clients);
close(epoll_fd);
return 1;
}
}
struct epoll_event events[MAX_EVENTS];
double start_time = get_time_double();
double end_time = start_time + TEST_DURATION_S;
while (get_time_double() < end_time) {
int n = epoll_wait(epoll_fd, events, MAX_EVENTS, 10); // 10ms timeout
double now = get_time_double();
for (int i = 0; i < n; ++i) {
Client* client = (Client*)events[i].data.ptr;
if (events[i].events & (EPOLLERR | EPOLLHUP)) {
close_client(client);
continue;
}
if (client->state == STATE_CONNECTING && (events[i].events & EPOLLOUT)) {
int result;
socklen_t result_len = sizeof(result);
getsockopt(client->fd, SOL_SOCKET, SO_ERROR, &result, &result_len);
if (result == 0) {
send_handshake(client);
} else {
close_client(client);
}
} else {
if (events[i].events & EPOLLIN) handle_read(client);
if (events[i].events & EPOLLOUT) handle_write(client);
}
}
// Publisher logic
if (all_subscribed) {
for (int i = NUM_SUBSCRIBERS; i < TOTAL_CLIENTS; ++i) {
Client* client = &clients[i];
if (client->state == STATE_RUNNING && client->write_len == 0 && now >= client->next_send_time) {
const char* channel = CHANNELS[rand() % NUM_CHANNELS];
char message[256];
int msg_len = snprintf(message, sizeof(message), "%.6f:Hello from publisher %d on channel %s", now, i - NUM_SUBSCRIBERS, channel);
char pub_msg[384];
int pub_msg_len = snprintf(pub_msg, sizeof(pub_msg), "pub %s %s", channel, message);
client->write_len = create_ws_frame(pub_msg, pub_msg_len, client->write_buf);
messages_sent++;
client->next_send_time = now + (1.0 / MESSAGES_PER_SECOND_PER_PUBLISHER);
epoll_ctl_mod(client->fd, EPOLLOUT | EPOLLIN | EPOLLET, client);
}
}
}
}
printf("\nTest duration finished. Shutting down clients...\n");
for (int i = 0; i < TOTAL_CLIENTS; ++i) {
close_client(&clients[i]);
}
print_report();
free(clients);
free(latencies);
close(epoll_fd);
return 0;
}
Reference in New Issue Copy Permalink