258 lines
7.1 KiB
C
258 lines
7.1 KiB
C
// SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause)
|
|
/* Copyright (c) 2020 Facebook */
|
|
#include <stdio.h>
|
|
#include <unistd.h>
|
|
#include <sys/resource.h>
|
|
#include <bpf/libbpf.h>
|
|
#include "replica.skel.h"
|
|
#include <sys/socket.h>
|
|
#include <netinet/in.h>
|
|
#include <arpa/inet.h>
|
|
#include <fcntl.h>
|
|
|
|
#include "replica.h"
|
|
|
|
|
|
struct cmd_node {
|
|
uint8_t *cmd;
|
|
size_t len;
|
|
struct cmd_node *next;
|
|
};
|
|
|
|
struct pending_queue {
|
|
struct cmd_node *head;
|
|
struct cmd_node *tail;
|
|
int count;
|
|
};
|
|
|
|
static void queue_init(struct pending_queue *q) {
|
|
q->head = q->tail = NULL;
|
|
q->count = 0;
|
|
}
|
|
|
|
static void queue_push(struct pending_queue *q, const uint8_t *cmd, size_t len) {
|
|
struct cmd_node *node = malloc(sizeof(*node));
|
|
if (!node) return;
|
|
node->cmd = malloc(len);
|
|
if (!node->cmd) { free(node); return; }
|
|
memcpy(node->cmd, cmd, len);
|
|
node->len = len;
|
|
node->next = NULL;
|
|
if (q->tail) q->tail->next = node;
|
|
else q->head = node;
|
|
q->tail = node;
|
|
q->count++;
|
|
}
|
|
|
|
static void queue_send_and_clear(struct pending_queue *q, int sock) {
|
|
struct cmd_node *node = q->head;
|
|
int sent_count = 0;
|
|
while (node) {
|
|
if (send(sock, node->cmd, node->len, 0) > 0) {
|
|
sent_count++;
|
|
}
|
|
struct cmd_node *tmp = node;
|
|
node = node->next;
|
|
free(tmp->cmd);
|
|
free(tmp);
|
|
}
|
|
if (sent_count > 0) {
|
|
printf("[QUEUE] Sent %d commands to slave\n", sent_count);
|
|
}
|
|
queue_init(q);
|
|
}
|
|
|
|
static void queue_free(struct pending_queue *q) {
|
|
struct cmd_node *node = q->head;
|
|
while (node) {
|
|
struct cmd_node *tmp = node;
|
|
node = node->next;
|
|
free(tmp->cmd);
|
|
free(tmp);
|
|
}
|
|
queue_init(q);
|
|
}
|
|
|
|
static int send_file(int sock, const char *path) {
|
|
FILE *fp = fopen(path, "rb");
|
|
if (!fp) {
|
|
printf("[ERROR] Failed to open file: %s\n", path);
|
|
return -1;
|
|
}
|
|
|
|
char buf[4096];
|
|
size_t n, total = 0;
|
|
while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
|
|
if (send(sock, buf, n, 0) < 0) {
|
|
fclose(fp);
|
|
printf("[ERROR] Failed to send file: %s (sent %zu bytes)\n", path, total);
|
|
return -1;
|
|
}
|
|
total += n;
|
|
}
|
|
fclose(fp);
|
|
printf("[FILE] Sent %s (%zu bytes)\n", path, total);
|
|
return 0;
|
|
}
|
|
|
|
// 全局状态（单 Slave 简化）
|
|
static enum state current_state = NOSLAVE;
|
|
static char slave_ip[16] = {0};
|
|
static int slave_port = 0;
|
|
static char array_file[128] = {0};
|
|
static char rbtree_file[128] = {0};
|
|
static char hash_file[128] = {0};
|
|
static struct pending_queue pending;
|
|
static int slave_sock = -1; // 连接 Slave 的 socket
|
|
|
|
// 连接 Slave
|
|
static int connect_slave() {
|
|
int sock = socket(AF_INET, SOCK_STREAM, 0);
|
|
if (sock < 0) return -1;
|
|
struct sockaddr_in addr;
|
|
memset(&addr, 0, sizeof(addr));
|
|
addr.sin_family = AF_INET;
|
|
addr.sin_port = htons(slave_port);
|
|
inet_pton(AF_INET, slave_ip, &addr.sin_addr);
|
|
if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
|
|
close(sock);
|
|
return -1;
|
|
}
|
|
return sock;
|
|
}
|
|
|
|
static void handle_event(void *ctx, int cpu, void *data, __u32 size) {
|
|
struct event *ev = (struct event *)data;
|
|
|
|
switch (ev->type) {
|
|
case EVENT_CREATE_SNAPSHOT_ASYNC:
|
|
printf("[EVENT] Type: CREATE_SNAPSHOT_ASYNC\n");
|
|
printf("[EVENT] Slave IP: %s, Port: %u\n", ev->data.sync.ip, ev->data.sync.port);
|
|
|
|
if (current_state == NOSLAVE) {
|
|
current_state = START;
|
|
strncpy(slave_ip, ev->data.sync.ip, sizeof(slave_ip));
|
|
slave_port = ev->data.sync.port;
|
|
queue_init(&pending);
|
|
slave_sock = connect_slave(); // 连接 Slave
|
|
if (slave_sock < 0) {
|
|
printf("Failed to connect to Slave %s:%d\n", slave_ip, slave_port);
|
|
current_state = NOSLAVE;
|
|
}
|
|
}
|
|
break;
|
|
case EVENT_COMPLETED_CMD:
|
|
printf("[EVENT] Type: COMPLETED_CMD\n");
|
|
printf("[EVENT] Command length: %llu bytes\n", ev->data.cmd.len);
|
|
|
|
if (current_state != NOSLAVE) {
|
|
queue_push(&pending, ev->data.cmd.cmd, ev->data.cmd.len);
|
|
}
|
|
break;
|
|
case EVENT_CREATE_SNAPSHOT_OK:
|
|
printf("[EVENT] Type: CREATE_SNAPSHOT_OK\n");
|
|
printf("[EVENT] Array file: %s\n", ev->data.ok.array_file);
|
|
printf("[EVENT] RBTree file: %s\n", ev->data.ok.rbtree_file);
|
|
printf("[EVENT] Hash file: %s\n", ev->data.ok.hash_file);
|
|
|
|
if (current_state == START) {
|
|
current_state = DONE;
|
|
strncpy(array_file, ev->data.ok.array_file, sizeof(array_file));
|
|
strncpy(rbtree_file, ev->data.ok.rbtree_file, sizeof(rbtree_file));
|
|
strncpy(hash_file, ev->data.ok.hash_file, sizeof(hash_file));
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
|
|
|
|
static void lost_event(void *ctx, int cpu, __u64 cnt) {
|
|
printf("Lost %llu events\n", cnt);
|
|
}
|
|
|
|
static int libbpf_print_fn(enum libbpf_print_level level, const char *format, va_list args)
|
|
{
|
|
return vfprintf(stderr, format, args);
|
|
}
|
|
|
|
int main(int argc, char **argv)
|
|
{
|
|
struct replica_bpf *skel;
|
|
int err;
|
|
|
|
/* Set up libbpf errors and debug info callback */
|
|
libbpf_set_print(libbpf_print_fn);
|
|
|
|
/* Open BPF application */
|
|
skel = replica_bpf__open();
|
|
if (!skel) {
|
|
fprintf(stderr, "Failed to open BPF skeleton\n");
|
|
return 1;
|
|
}
|
|
|
|
/* ensure BPF program only handles write() syscalls from our process */
|
|
skel->bss->my_pid = getpid();
|
|
|
|
/* Load & verify BPF programs */
|
|
err = replica_bpf__load(skel);
|
|
if (err) {
|
|
fprintf(stderr, "Failed to load and verify BPF skeleton\n");
|
|
goto cleanup;
|
|
}
|
|
|
|
/* Attach tracepoint handler */
|
|
err = replica_bpf__attach(skel);
|
|
if (err) {
|
|
fprintf(stderr, "Failed to attach BPF skeleton\n");
|
|
goto cleanup;
|
|
}
|
|
|
|
printf("Successfully started! Please run `sudo cat /sys/kernel/debug/tracing/trace_pipe` "
|
|
"to see output of the BPF programs.\n");
|
|
|
|
|
|
struct perf_buffer *pb = perf_buffer__new(bpf_map__fd(skel->maps.channel), 8, handle_event, lost_event, NULL, NULL);
|
|
if(!pb){
|
|
goto cleanup;
|
|
}
|
|
|
|
#if 0
|
|
while(1){
|
|
perf_buffer__poll(pb, 1000);
|
|
}
|
|
#else
|
|
while (1) {
|
|
perf_buffer__poll(pb, 1000); // 处理事件
|
|
|
|
// 循环中检查状态并发送
|
|
if (current_state == DONE && slave_sock >= 0) {
|
|
// 发送快照文件
|
|
if (send_file(slave_sock, array_file) == 0 &&
|
|
send_file(slave_sock, rbtree_file) == 0 &&
|
|
send_file(slave_sock, hash_file) == 0) {
|
|
current_state = ONLINE;
|
|
printf("Snapshot sent, state to ONLINE\n");
|
|
} else {
|
|
printf("Failed to send snapshot\n");
|
|
current_state = NOSLAVE;
|
|
close(slave_sock);
|
|
slave_sock = -1;
|
|
}
|
|
}
|
|
|
|
if (current_state == ONLINE && slave_sock >= 0) {
|
|
// 发送 pending
|
|
queue_send_and_clear(&pending, slave_sock);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
perf_buffer__free(pb);
|
|
|
|
cleanup:
|
|
queue_free(&pending);
|
|
if (slave_sock >= 0) close(slave_sock);
|
|
replica_bpf__destroy(skel);
|
|
return -err;
|
|
}
|