Add udp heartbeat mux example

2013-01-21 23:11:36 +01:00
parent b378b205c9
commit de9eab51a5
1 changed files with 292 additions and 0 deletions
@@ -0,0 +1,292 @@
+#define DEBUG
+#include <getopt.h>
+#include <errno.h>
+#include <stdlib.h>
+#include <time.h>
+#include <unistd.h>
+#include <unistd.h>
+#include <stdio.h>
+#include <string.h>
+#include <stdint.h>
+#include <sys/socket.h>
+#include <sys/queue.h>
+#include <netinet/in.h>
+#include <arpa/inet.h>
+#include <ucore/ucore_logging.h>
+#include <ucore/ucore_logging.h>
+#include <ucore/iomux.h>
+#include <ucore/ucore_utils.h>
+#include <ucore/ucore_pack.h>
+
+#define HB_MAGIX 0xBEA7BEA7
+
+#define LMAIN 0
+#define PEER_HB_SEC 5
+#define PEER_HB_MISS 3
+
+static struct uc_log_module modules[] = {
+    {
+        .id = LMAIN,
+        .short_name = "MAIN",
+        .long_name = "Main logging module",
+        .log_level = UC_LL_DEBUG
+    }
+};
+
+static struct uc_log_modules log_modules = {
+    .mods = modules,
+    .cnt = ARRAY_SIZE(modules)
+};
+
+enum PeerState {
+    PeerUnknown,
+    PeerAlive,
+    PeerDead
+};
+
+enum PeerEvent {
+    PeerReceived,
+    PeerTimedout,
+};
+
+struct HBPeer {
+    struct UCTimer timer;
+    struct sockaddr_in peer_addr;
+    time_t last_ok;
+    enum PeerState state;
+    TAILQ_ENTRY(HBPeer) entry;
+};
+
+struct HBContext {
+    struct IOMuxFD sock_fd;
+    TAILQ_HEAD(, HBPeer) peers;
+    struct IOMux *mux;
+};
+
+struct HBPeer *hb_peer_find(struct HBContext *ctx, const struct sockaddr_in *peer_addr);
+void peer_send(struct HBContext *ctx, struct HBPeer *peer);
+void peer_do_state(struct HBPeer *peer, enum PeerEvent event);
+
+
+void peer_timer_cb(struct UCTimers *timers, struct UCTimer *timer)
+{
+    struct HBPeer *peer = CONTAINER_OF(timer, struct HBPeer, timer);
+    struct HBContext *ctx = timer->cookie_ptr;
+    time_t now;
+
+    uc_timer_add(timers, timer, 5 , 0);
+    now = time(NULL);
+
+    if (now - peer->last_ok > PEER_HB_MISS * PEER_HB_SEC) {
+        peer_do_state(peer, PeerTimedout);
+    }
+
+    peer_send(ctx, peer);
+}
+
+void peer_add_addr(struct HBContext *ctx, const struct sockaddr_in *addr)
+{
+    struct HBPeer *peer = hb_peer_find(ctx, addr);
+    if (peer != NULL) {
+        UC_LOGF(UC_LL_INFO, LMAIN, "Ignorming adding existing peer %s:%u\n",
+            inet_ntoa(addr->sin_addr), addr->sin_port);
+        return;
+    }
+
+    peer = calloc(1, sizeof *peer);
+    peer->peer_addr = *addr;
+    TAILQ_INSERT_TAIL(&ctx->peers, peer, entry);
+
+    peer->timer.callback = peer_timer_cb;
+    peer->timer.cookie_ptr = ctx;
+
+    uc_timer_add(iomux_get_timers(ctx->mux), &peer->timer, PEER_HB_SEC , 0);
+}
+
+void peer_add(struct HBContext *ctx, const char *host, uint16_t port)
+{
+    struct sockaddr_in peer_addr;
+
+    UC_DEBUGF(LMAIN, "Adding peer %s:%u\n", host, port);
+    memset(&peer_addr, 0, sizeof peer_addr);
+
+    peer_addr.sin_family = AF_INET;
+    peer_addr.sin_addr.s_addr = inet_addr(host);
+    peer_addr.sin_port = htons(port);
+
+    peer_add_addr(ctx, &peer_addr);
+}
+
+void peer_send(struct HBContext *ctx, struct HBPeer *peer)
+{
+    uint8_t buf[4];
+    ssize_t rc;
+
+    uc_pack_32_be(HB_MAGIX, buf);
+
+    UC_DEBUGF(LMAIN, "Sending to %s:%u\n", inet_ntoa(peer->peer_addr.sin_addr), 
+        ntohs(peer->peer_addr.sin_port));
+
+    rc = sendto(ctx->sock_fd.fd, buf, sizeof buf, 0,
+            (struct sockaddr*)&peer->peer_addr, sizeof peer->peer_addr);
+    if (rc == -1) {
+        UC_DEBUGF(LMAIN, "Could not send packet: %s\n", strerror(errno));
+    }
+}
+
+struct HBPeer *hb_peer_find(struct HBContext *ctx, const struct sockaddr_in *peer_addr)
+{
+    struct HBPeer *peer;
+
+    TAILQ_FOREACH(peer, &ctx->peers, entry) {
+        if (peer_addr->sin_addr.s_addr == peer->peer_addr.sin_addr.s_addr 
+            && peer_addr->sin_port == peer->peer_addr.sin_port) {
+            break;
+        }
+    }
+
+    return peer;
+}
+
+void peer_do_state(struct HBPeer *peer, enum PeerEvent event)
+{
+    if (peer->state == PeerUnknown || peer->state == PeerDead) {
+        if (event == PeerReceived) {
+            UC_LOGF(UC_LL_INFO, LMAIN, "Peer is now alive\n");
+            peer->state = PeerAlive;
+        }
+    } else if (peer->state == PeerAlive) {
+        if (event == PeerTimedout) {
+            UC_LOGF(UC_LL_INFO, LMAIN, "Peer is dead !\n");
+            peer->state = PeerDead;
+        }
+    }
+
+    if (peer->state == PeerUnknown && event == PeerTimedout) {
+        peer->state = PeerDead;
+    }
+}
+
+void peer_received(struct HBPeer *peer)
+{
+    peer->last_ok = time(NULL);
+    peer_do_state(peer, PeerReceived);
+}
+
+void  hb_read(struct IOMux *mux, struct IOMuxFD *fd, unsigned int event)
+{
+    struct HBContext *ctx = CONTAINER_OF(fd, struct HBContext, sock_fd);
+    ssize_t rc;
+    struct sockaddr_in peer_addr;
+    socklen_t peer_len = sizeof peer_addr;
+    uint8_t buf[128];
+    uint32_t magic;
+    struct HBPeer *peer;
+
+    rc = recvfrom(fd->fd, buf, sizeof buf, 0, (struct sockaddr*)&peer_addr, &peer_len);
+    if (rc == -1 && errno == EWOULDBLOCK) {
+        UC_DEBUGF(LMAIN, "recvfrom EWOULDBLOCK\n");
+        return;
+    }
+
+    UC_LOGF(UC_LL_INFO, LMAIN, "Message from %s:%u\n", inet_ntoa(peer_addr.sin_addr),
+        ntohs(peer_addr.sin_port));
+
+    if (rc < 4) {
+        UC_DEBUGF(LMAIN, "short read of %u\n", (unsigned)rc);
+        return;
+    }
+
+    magic = uc_unpack_32_be(buf);
+    if (magic != HB_MAGIX) {
+        UC_LOGF(UC_LL_WARNING, LMAIN, "Received unknown message, magic = 0x%X\n", magic);
+        return;
+    }
+
+    peer = hb_peer_find(ctx, &peer_addr);
+    if (peer == NULL) {
+        UC_LOGF(UC_LL_WARNING, LMAIN, "Found no peer\n");
+        return;
+    }
+
+    peer_received(peer);
+    
+    peer_send(ctx, peer);
+}
+
+void init_hb_context(struct HBContext *ctx, uint16_t local_port)
+{
+    int sock_fd;
+    struct sockaddr_in addr;
+    
+    memset(&addr, 0, sizeof addr);
+    addr.sin_family = AF_INET;
+    addr.sin_port = htons(local_port);
+    addr.sin_addr.s_addr = INADDR_ANY;
+
+    sock_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+    if (sock_fd == -1) {
+        UC_LOGF(UC_LL_ERROR, LMAIN, "socket failed : %s\n", strerror(errno));
+        return;
+
+    }
+
+    if (bind(sock_fd, (struct sockaddr*)&addr, sizeof addr) != 0) {
+        UC_LOGF(UC_LL_ERROR, LMAIN,  "Cannot bind socket to port %u : %s\n", 
+            local_port, strerror(errno));
+        close(sock_fd);
+        return;
+    }
+
+    TAILQ_INIT(&ctx->peers);
+
+    ctx->sock_fd.fd = sock_fd;
+    ctx->sock_fd.what = MUX_EV_READ;
+    ctx->sock_fd.callback = hb_read;
+}
+
+int main(int argc, char *argv[])
+{
+    struct HBContext ctx;
+    struct IOMux *mux;
+    uint16_t local_port = 23456;
+    uint16_t remote_port = 23456;
+    int c;
+
+    uc_log_init(&log_modules);
+    uc_log_add_destination(uc_log_new_stderr(UC_LL_DEBUG, 1));
+
+    mux = iomux_create(IOMUX_TYPE_DEFAULT);
+    ctx.mux = mux;
+
+    init_hb_context(&ctx, local_port);
+
+    while ((c = getopt(argc, argv, "p:P:a:")) != -1) {
+        switch (c) {
+        case 'p':
+            local_port = atoi(optarg);
+            break;
+        case 'P':
+            remote_port = atoi(optarg);
+            break;
+        case 'a':
+            peer_add(&ctx, optarg, remote_port);
+            break;
+        default:
+            printf("Unknown option -%c\n", c);
+            return 1;
+        }
+    }
+
+
+
+
+    UC_DEBUGF(LMAIN, "Creating IOMux\n");
+
+    iomux_register_fd(mux, &ctx.sock_fd);
+
+    iomux_run(mux);
+
+
+    return 0; 
+}