基于 UDP 的 golang P2P 实例

好风 发表于 2017-01-09T02:22:13.544049Z
引用地址:https://plus.ooclab.com/note/article/1337

UDP 是 P2P 架构基础

代码

server

package main

import (
    "fmt"
    "net"
    "os"
    "strings"
    "sync"
    "time"
)

type Client struct {
    Addr       net.Addr
    lastActive time.Time
}

type ClientPool struct {
    p map[string]*Client
    m sync.Mutex
}

func (c ClientPool) Add(addr net.Addr) {
    c.m.Lock()
    c.p[addr.String()] = &Client{
        Addr:       addr,
        lastActive: time.Now(),
    }
    c.m.Unlock()
}

func (c ClientPool) Friends(addr net.Addr) string {
    friends := make([]string, 1)
    i := addr.String()
    for k := range c.p {
        if k != i {
            friends = append(friends, k)
        }
    }
    return strings.Join(friends, ",")
}

func (c ClientPool) PrintAll() {
    for k := range c.p {
        fmt.Println("--", k)
    }
}

func (c ClientPool) Clean() {
    c.m.Lock()
    defer c.m.Unlock()
    for k, client := range c.p {
        if time.Since(client.lastActive) > 5*time.Second {
            delete(c.p, k)
            fmt.Println("delete", k)
        }
    }
}

func main() {
    if len(os.Args) != 2 {
        fmt.Printf("Usage: %s IP:Port\n", os.Args[0])
        os.Exit(1)
    }

    addr := os.Args[1]

    laddr, err := net.ResolveUDPAddr("udp", addr)
    if err != nil {
        fmt.Printf("resolve udp addr (%s) failed: %s\n", addr, err)
        os.Exit(1)
    }

    conn, err := net.ListenUDP("udp", laddr)
    if err != nil {
        fmt.Printf("listen udp failed: %s\n", err)
        os.Exit(1)
    }
    cp := &ClientPool{
        p: map[string]*Client{},
    }
    go func() {
        for {
            cp.Clean()
            time.Sleep(3 * time.Second)
        }
    }()
    buf := make([]byte, 8800)
    for {
        _, raddr, _ := conn.ReadFromUDP(buf)
        // fmt.Println("n, raddr, err = ", n, raddr, err)
        // fmt.Println("\n" + hex.Dump(buf[0:n]))
        cp.Add(raddr)
        friends := cp.Friends(raddr)
        n, err := conn.WriteToUDP([]byte(friends), raddr)
        if err != nil {
            fmt.Println("Error: ", err)
        }
        cp.PrintAll()
        fmt.Printf("send to %s: friends = %s, n = %d\n", raddr.String(), friends, n)
    }
}

client

package main

import (
    "fmt"
    "net"
    "os"
    "strings"
    "time"
)

func usageCheck() {
    if len(os.Args) != 2 {
        fmt.Printf("Usage: %s IP:Port\n", os.Args[0])
        os.Exit(1)
    }
}

func errCheck(err error) {
    if err != nil {
        fmt.Println("panic: ", err)
        os.Exit(1)
    }
}

var friendCh = make(chan []string, 1)

func getFriends(conn *net.UDPConn, raddr *net.UDPAddr) []string {
    for {
        _, err := conn.WriteToUDP([]byte("GET /friends"), raddr)
        errCheck(err)
        select {
        case friends := <-friendCh:
            return friends
        case <-time.After(time.Second * 5):
            fmt.Println("wait friends timeout ...")
        }
    }
}

func main() {
    usageCheck()

    raddr, err := net.ResolveUDPAddr("udp", os.Args[1])
    errCheck(err)
    conn, err := net.ListenUDP("udp", nil)
    errCheck(err)

    go func() {
        for {
            friends := getFriends(conn, raddr)
            // fmt.Println("get friends: ", friends)

            for _, friend := range friends {
                if len(friend) == 0 {
                    continue
                }
                faddr, err := net.ResolveUDPAddr("udp", friend)
                if err != nil {
                    fmt.Printf("resolve udp addr %s failed: %s\n", friend, err)
                    continue
                }

                buf := []byte("Hello, " + friend)
                n, err := conn.WriteToUDP(buf, faddr)
                if err != nil {
                    fmt.Println("Error: n, err = ", n, err)
                }
                // fmt.Printf("write to %s success: %d\n", faddr, n)
            }
            // fmt.Println("sleep 3 seconds ...")
            time.Sleep(3 * time.Second)
        }
    }()

    buf := make([]byte, 8800)
    for {
        n, raddr, err := conn.ReadFromUDP(buf)
        if err != nil {
            fmt.Println("Error: n, raddr, err = ", n, raddr, err)
            continue
        }
        msg := string(buf[0:n])
        if strings.HasPrefix(msg, "Hello") {
            fmt.Printf("RECV from %s: %s\n", raddr, msg)
        } else {
            // friends
            friendCh <- strings.Split(msg, ",")
        }
    }

}