[ISSUE] Handling SIGINT and SIGTERM for Graceful Shutdown in Go Websocket App

[matterai]

Is there an existing issue for this?

• I have searched the existing issues

Current Behavior

Hey!

I've got a main.go file where I'm playing around with this code snippet:

ctx, cancelCtx := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer cancelCtx()

service := NewService(ctx)
service.Connect()

Inside the Run() method, I'm setting up a connection with DialContext(s.context, url, nil) and diving into an endless loop that listens for signals. My goal here is to catch SIGINT and SIGTERM so I can neatly wrap up the connection and close the app on a good note. Here's what the service code looks like:

func NewService(context context.Context) *GatewayService {
	return &GatewayService{
		context: context,
	}
}

func (s *Service) Connect() error {
	url := "wss://dial.ws"
	conn, _, err := websocket.DefaultDialer.DialContext(s.context, url, nil)
	if err != nil {
		log.Printf("Error dialing: %v", err)
		return err
	}

	for {
		select {
		case <-s.context.Done():
			log.Println("Context done")
			if err := s.conn.Close(); err != nil {
				log.Printf("Error closing connection: %v", err)
				return err
			}

			return s.context.Err()

		default:
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Printf("Error reading message: %v", err)
				return err
			}

			log.Println(string(message))
		}
	}
}

When I give it a go and trigger a SIGTERM, the app just keeps on running. It looks like conn.ReadMessage() is holding up the main go-routine while it waits for new bytes from the websocket channel. I might be missing something here, but this seems a bit off to me. Is this how it's supposed to work? What's the best way to deal with termination signals with this library?

Thanks a bunch for your insights!

Expected Behavior

No response

Steps To Reproduce

No response

Anything else?

No response

[GreenMarmot]

Use two goroutines, one to read the connection and the other to handle signal. Here's an example.

[matterai]

Yep, I fixed it yesterday already. I made a goroutine for reading messages. The main goroutine is awaiting signals in for-cycle in the end of my Connect func:

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				s.logger.Printf("Error reading message: %v", err)
				return
			}

			_, err = s.handler.HandleMessage(message)
			if err != nil {
				s.logger.Printf("Error handling message: %v", err)
				return
			}
		}
	}()

	for {
		select {
		case <-s.context.Done():
			s.logger.Println("Conext done, closing connection")
			return s.context.Err()
		case <-done:
			s.logger.Println("Message handler done, closing connection")
			return s.context.Err()
		}
	}

Thank you. Sorry for bothering.