Изменить обработку прерываний, чтобы закрыть канал при прерывании.Это позволяет нескольким процедурам ожидать события, ожидая закрытия канала.
shutdown := make(chan struct{})
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, os.Interrupt)
go func() {
<-interrupt
log.Println("interrupt")
close(shutdown)
}()
Перемещение кода для соединения в функцию.Этот код является копией и вставкой из вопроса с двумя изменениями: канал прерывания заменяется каналом отключения;Функция уведомляет sync.WaitGroup о завершении функции.
func connect(u string, shutdown chan struct{}, wg *sync.WaitGroup) {
defer wg.Done()
log.Printf("connecting to %s", u)
c, _, err := websocket.DefaultDialer.Dial(u, nil)
if err != nil {
log.Fatal("dial:", err)
}
defer c.Close()
done := make(chan struct{})
go func() {
defer close(done)
for {
_, message, err := c.ReadMessage()
if err != nil {
log.Println("read:", err)
return
}
log.Printf("recv: %s", message)
}
}()
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
for {
select {
case <-done:
return
case t := <-ticker.C:
err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
if err != nil {
log.Println("write:", err)
return
}
case <-shutdown:
// Cleanly close the connection by sending a close message and then
// waiting (with timeout) for the server to close the connection.
err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil {
log.Println("write close:", err)
return
}
select {
case <-done:
case <-time.After(time.Second):
}
return
}
}
}
Объявите sync.WaitGroup в main().Для каждой конечной точки веб-сокета, к которой вы хотите подключиться, увеличьте WaitGroup и запустите процедуру подключения этой конечной точки.После запуска goroutines, дождитесь завершения WaitGroup для goroutines.
var wg sync.WaitGroup
for _, u := range endpoints { // endpoints is []string
// where elements are URLs
// of endpoints to connect to.
wg.Add(1)
go connect(u, shutdown, &wg)
}
wg.Wait()
Приведенный выше код с правкой, позволяющей запустить его на эхо-сервере Gorilla, размещен на детской площадке .