Go : Cancel Context inside a For loop

Question

I am trying to create a UDP server in Golang to Listen at a port for eg. 1234. I have a client which sends the start/stop message to this server.

On receiving of message "start", the server will start sending random data to this client and on the stop, the server will stop sending to the client.

For this purpose, I am using context to create a goroutine to send the data and cancel it when it gets "stop".

The error I am getting is the program works fine for one client, but if I start the client again the data is not sent again.

Any help would be appreciated?

UDP server Code:

package main

import (
    "context"
    "fmt"
    "math/rand"
    "net"
    "time"
)

func generateMessageToUDP(ctx context.Context, addr *net.UDPAddr) {
    // stop writing to UDP
    done := false
    fmt.Println("Generating message to UDP client", addr)
    conn, err := net.DialUDP("udp", nil, addr)
    if err != nil {
        fmt.Println("Error: ", err)
    }
    defer func(conn *net.UDPConn) {
        err := conn.Close()
        if err != nil {
            fmt.Println("Error in closing the UDP Connection: ", err)
        }
    }(conn)
    // write to address using UDP connection
    go func() {
        for i := 0; !done; i++ {
            RandomInt := rand.Intn(100)
            fmt.Println("Random Int: ", RandomInt)
            _, err = conn.Write([]byte(fmt.Sprintf("%d", RandomInt)))
            fmt.Println("Sent ", RandomInt, " to ", addr)
            time.Sleep(time.Second * 1)
        }
    }()
    <-ctx.Done()
    fmt.Println("Stopping writing to UDP client", addr)
    done = true
}

//var addr *net.UDPAddr
//var conn *net.UDPConn

func main() {
    fmt.Println("Hi this is a UDP server")
    udpServer, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(0, 0, 0, 0), Port: 5010})
    if err != nil {
        fmt.Println("Error: ", err)
    }
    defer func(udpServer *net.UDPConn) {
        err := udpServer.Close()
        if err != nil {
            fmt.Println("Error in closing the UDP Connection: ", err)
        }
    }(udpServer)
    // create a buffer to read data into
    buffer := make([]byte, 1024)
    ctx, cancel := context.WithCancel(context.Background())
    for {
        // read the incoming connection into the buffer
        n, addr, err := udpServer.ReadFromUDP(buffer)
        fmt.Println("Recieved ", string(buffer[0:n]), " from ", addr)
        if err != nil {
            fmt.Println("Error: ", err)
        }
        fmt.Println("Received ", string(buffer[0:n]), " from ", addr)
        if string(buffer[0:n]) == "stop" {
            fmt.Println("Stopped listening")
            cancel()
            continue
        } else if string(buffer[0:n]) == "start" {
            // send a response back to the client
            _, err = udpServer.WriteToUDP([]byte("Hi, I am a UDP server"), addr)
            if err != nil {
                fmt.Println("Error: ", err)
            }
            // start a routine to generate messages to the client
            generateMessageToUDP(ctx, addr)
        } else {
            fmt.Println("Unknown command")
        }
    }
}

Client Code:

package main

import (
    "fmt"
    "net"
    "time"
)

func main() {
    fmt.Println("Hello, I am a client")

    // Create a new client
    localAddr, err := net.ResolveUDPAddr("udp", ":5011")
    client3, err := net.DialUDP("udp", localAddr, &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 5010})
    if err != nil {
        fmt.Println(err)
        return
    }
    defer client3.Close()
    _, err = client3.Write([]byte("start"))
    if err != nil {
        fmt.Println(err)
        return
    }
    fmt.Println("Message sent. Sleeping for 5 seconds")
    time.Sleep(time.Second * 5)
    fmt.Println("Sending stop message")
    _, err = client3.Write([]byte("stop"))
    if err != nil {
        fmt.Println(err)
    }
}

Answer 1

Okay, I did a simple hack on the server and added a label Start before creating a context and when I cancel the context, I addded goto label. This means when the task get cancelled it will again create the context and start doings its job