ldb/kvs_slave.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>


static int recv_exact(int sockfd, char *buf, size_t len) {
    size_t received = 0;
    while (received < len) {
        ssize_t ret = recv(sockfd, buf + received, len - received, 0);
        if (ret <= 0) {
            if (ret == 0) {
                fprintf(stderr, "connection closed by peer\n");
            } else {
                perror("recv");
            }
            return -1;
        }
        received += ret;
    }
    return 0;
}

static int receive_snapshot_file(int sockfd, const char *filename) {
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        perror("fopen");
        return -1;
    }

    char buf[4096];
    ssize_t n;
    
    while ((n = recv(sockfd, buf, sizeof(buf), 0)) > 0) {
        if (fwrite(buf, 1, n, fp) != n) {
            perror("fwrite");
            fclose(fp);
            return -1;
        }
    }
    
    if (n < 0) {
        perror("recv");
        fclose(fp);
        return -1;
    }
    
    fclose(fp);
    return 0;
}


int slave_bootstrap(
    const char *listen_ip,
    int listen_port,
    const char *master_ip,
    int master_port
) {
    int listen_fd = -1;
    int master_fd = -1;
    int snapshot_fds[3] = {-1, -1, -1};
    int ret = -1;

    // 1. 创建监听socket
    listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd < 0) {
        perror("socket");
        goto cleanup;
    }

    // 设置SO_REUSEADDR
    int opt = 1;
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("setsockopt");
        goto cleanup;
    }

    struct sockaddr_in listen_addr;
    memset(&listen_addr, 0, sizeof(listen_addr));
    listen_addr.sin_family = AF_INET;
    listen_addr.sin_port = htons(listen_port);
    listen_addr.sin_addr.s_addr = inet_addr(listen_ip);

    if (bind(listen_fd, (struct sockaddr*)&listen_addr, sizeof(listen_addr)) < 0) {
        perror("bind");
        goto cleanup;
    }

    if (listen(listen_fd, 5) < 0) {
        perror("listen");
        goto cleanup;
    }

    printf("Slave listening on %s:%d\n", listen_ip, listen_port);


    // 2. 连接master并发送SSYNC命令
    master_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (master_fd < 0) {
        perror("socket");
        goto cleanup;
    }

    struct sockaddr_in master_addr;
    memset(&master_addr, 0, sizeof(master_addr));
    master_addr.sin_family = AF_INET;
    master_addr.sin_port = htons(master_port);
    master_addr.sin_addr.s_addr = inet_addr(master_ip);

    if (connect(master_fd, (struct sockaddr*)&master_addr, sizeof(master_addr)) < 0) {
        perror("connect to master");
        goto cleanup;
    }

    printf("Connected to master %s:%d\n", master_ip, master_port);

    // 构造RESP协议消息: SSYNC listen_ip listen_port
    char ssync_cmd[256];
    int cmd_len = snprintf(ssync_cmd, sizeof(ssync_cmd),
                          "*3\r\n$5\r\nSSYNC\r\n$%zu\r\n%s\r\n$%d\r\n%d\r\n",
                          strlen(listen_ip), listen_ip, 
                          (int)snprintf(NULL, 0, "%d", listen_port), listen_port);

    if (send(master_fd, ssync_cmd, cmd_len, 0) < 0) {
        perror("send SSYNC");
        goto cleanup;
    }

    printf("Sent SSYNC command to master\n");

    // 3. 接收master的+OK\r\n回包
    char resp[32];
    if (recv_exact(master_fd, resp, 5) < 0) {
        fprintf(stderr, "Failed to receive OK from master\n");
        goto cleanup;
    }
    
    if (memcmp(resp, "+OK\r\n", 5) != 0) {
        fprintf(stderr, "Unexpected response from master: %.5s\n", resp);
        goto cleanup;
    }

    printf("Received OK from master\n");


    // 4. Accept来自master的三个连接，接收三个snapshot文件
    const char *snapshot_files[3] = {
        "data/kvs_array.db",
        "data/kvs_rbtree.db",
        "data/kvs_hash.db"
    };

    for (int i = 0; i < 3; i++) {
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
        
        snapshot_fds[i] = accept(listen_fd, (struct sockaddr*)&client_addr, &addr_len);
        if (snapshot_fds[i] < 0) {
            perror("accept");
            goto cleanup;
        }

        printf("Accepted connection %d from %s:%d\n", 
               i + 1, inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));

        if (receive_snapshot_file(snapshot_fds[i], snapshot_files[i]) < 0) {
            fprintf(stderr, "Failed to receive snapshot file %d\n", i + 1);
            goto cleanup;
        }

        printf("Received snapshot file: %s\n", snapshot_files[i]);
        close(snapshot_fds[i]);
        snapshot_fds[i] = -1;
    }

    // 5. 关闭监听和snapshot连接，发送SREADY
    close(listen_fd);
    listen_fd = -1;

    const char *sready_cmd = "*1\r\n$6\r\nSREADY\r\n";
    if (send(master_fd, sready_cmd, strlen(sready_cmd), 0) < 0) {
        perror("send SREADY");
        goto cleanup;
    }

    printf("Sent SREADY to master\n");

    // 6. 接收回包+OK\r\n
    if (recv_exact(master_fd, resp, 5) < 0) {
        fprintf(stderr, "Failed to receive final OK from master\n");
        goto cleanup;
    }
    
    if (memcmp(resp, "+OK\r\n", 5) != 0) {
        fprintf(stderr, "Unexpected final response from master: %.5s\n", resp);
        goto cleanup;
    }

    printf("Received final OK from master, bootstrap complete\n");
    
    ret = 0;

cleanup:
    if (master_fd >= 0) close(master_fd);
    if (listen_fd >= 0) close(listen_fd);
    for (int i = 0; i < 3; i++) {
        if (snapshot_fds[i] >= 0) close(snapshot_fds[i]);
    }
    
    return ret;
}