فهرست منبع

add quit function

Joyit 3 هفته پیش
والد
کامیت
a16ea27d72
1فایلهای تغییر یافته به همراه22 افزوده شده و 0 حذف شده
  1. 22 0
      hub.go

+ 22 - 0
hub.go

@@ -35,6 +35,8 @@ func subStr(str string, length int) string {
 
 type Hub struct {
 	sync.Mutex
+	ctx      context.Context // 为了方便退出而建立
+	cancel   context.CancelFunc
 	cf       *config.Config
 	globalID uint16
 	channel  string       // 本地频道信息
@@ -474,6 +476,8 @@ func (h *Hub) GetWithStruct(gd *GetData, backFunc GetBackFunc) (count int) {
 			}
 		case <-timer.C:
 			return
+		case <-h.ctx.Done():
+			return
 		}
 	}
 }
@@ -729,6 +733,8 @@ func (h *Hub) ConnectToServer(channel string, force bool, host *HostInfo) (err e
 	case <-ctx.Done():
 		done = true
 		return errors.New("timeout")
+	case <-h.ctx.Done():
+		return errors.New("quit")
 	}
 
 	// 如果 host 是代理,将代理信息添加到channel中
@@ -873,10 +879,23 @@ func (h *Hub) checkConnect() {
 				}
 				return true
 			})
+		case <-h.ctx.Done():
+			return
 		}
 	}
 }
 
+// 退出所有的连接
+func (h *Hub) Quit() {
+	h.cancel()
+	h.lines.Range(func(id int, line *Line) bool {
+		if line.state == Connected {
+			line.Close(true)
+		}
+		return true
+	})
+}
+
 // 建立一个集线器
 // connectFunc 用于监听连接状态的函数,可以为nil
 func NewHub(
@@ -896,7 +915,10 @@ func NewHub(
 	if cf == nil {
 		cf = config.NewConfig()
 	}
+	ctx, cancel := context.WithCancel(context.Background())
 	h = &Hub{
+		ctx:               ctx,
+		cancel:            cancel,
 		cf:                cf,
 		globalID:          uint16(time.Now().UnixNano()) % config.ID_MAX,
 		channel:           channel,