executer_manager(engine): Support watching executor online/offline ev…

…ents (#5565)

ref #5513

engine/servermaster/executor_manager.go

@@ -66,6 +66,8 @@ type ExecutorManagerImpl struct {
 
 	rescMgr resource.RescMgr
 	logRL   *rate.Limiter
+
+	notifier *notifier.Notifier[model.ExecutorStatusChange]
 }
 
 // NewExecutorManagerImpl creates a new ExecutorManagerImpl instance
@@ -78,6 +80,7 @@ func NewExecutorManagerImpl(initHeartbeatTTL, keepAliveInterval time.Duration, c
 		keepAliveInterval: keepAliveInterval,
 		rescMgr:           resource.NewCapRescMgr(),
 		logRL:             rate.NewLimiter(rate.Every(time.Second*5), 1 /*burst*/),
+		notifier:          notifier.NewNotifier[model.ExecutorStatusChange](),
 	}
 }
 
@@ -99,6 +102,11 @@ func (e *ExecutorManagerImpl) removeExecutorImpl(id model.ExecutorID) error {
 			Time: time.Now(),
 		})
 	}
+
+	e.notifier.Notify(model.ExecutorStatusChange{
+		ID: id,
+		Tp: model.EventExecutorOffline,
+	})
 	return nil
 }
 
@@ -150,6 +158,10 @@ func (e *ExecutorManagerImpl) RegisterExec(info *model.NodeInfo) {
 	}
 	e.mu.Lock()
 	e.executors[info.ID] = exec
+	e.notifier.Notify(model.ExecutorStatusChange{
+		ID: info.ID,
+		Tp: model.EventExecutorOnline,
+	})
 	e.mu.Unlock()
 	e.rescMgr.Register(exec.ID, exec.Addr, model.RescUnit(exec.Capability))
 }
@@ -293,7 +305,18 @@ func (e *ExecutorManagerImpl) GetAddr(executorID model.ExecutorID) (string, bool
 // WatchExecutors implements the ExecutorManager interface.
 func (e *ExecutorManagerImpl) WatchExecutors(
 	ctx context.Context,
-) ([]model.ExecutorID, *notifier.Receiver[model.ExecutorStatusChange], error) {
-	// TODO This method will be implemented before we enable local file GC.
-	panic("implement me")
+) (snap []model.ExecutorID, receiver *notifier.Receiver[model.ExecutorStatusChange], err error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	for executorID := range e.executors {
+		snap = append(snap, executorID)
+	}
+
+	if err := e.notifier.Flush(ctx); err != nil {
+		return nil, nil, err
+	}
+
+	receiver = e.notifier.NewReceiver()
+	return
 }

engine/servermaster/executor_manager_test.go

@@ -78,3 +78,71 @@ func TestExecutorManager(t *testing.T) {
 	require.NotNil(t, resp.Err)
 	require.Equal(t, pb.ErrorCode_UnknownExecutor, resp.Err.GetCode())
 }
+
+func TestExecutorManagerWatch(t *testing.T) {
+	t.Parallel()
+
+	heartbeatTTL := time.Millisecond * 400
+	checkInterval := time.Millisecond * 10
+	mgr := NewExecutorManagerImpl(heartbeatTTL, checkInterval, nil)
+	mgr.Start(context.Background())
+
+	// register an executor server
+	executorAddr := "127.0.0.1:10001"
+	registerReq := &pb.RegisterExecutorRequest{
+		Address:    executorAddr,
+		Capability: 2,
+	}
+	info, err := mgr.AllocateNewExec(registerReq)
+	require.Nil(t, err)
+
+	executorID1 := info.ID
+	snap, stream, err := mgr.WatchExecutors(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, []model.ExecutorID{executorID1}, snap)
+
+	// register another executor server
+	executorAddr = "127.0.0.1:10002"
+	registerReq = &pb.RegisterExecutorRequest{
+		Address:    executorAddr,
+		Capability: 2,
+	}
+	info, err = mgr.AllocateNewExec(registerReq)
+	require.Nil(t, err)
+
+	executorID2 := info.ID
+	event := <-stream.C
+	require.Equal(t, model.ExecutorStatusChange{
+		ID: executorID2,
+		Tp: model.EventExecutorOnline,
+	}, event)
+
+	newHeartbeatReq := func(executorID model.ExecutorID) *pb.HeartbeatRequest {
+		return &pb.HeartbeatRequest{
+			ExecutorId: string(executorID),
+			Status:     int32(model.Running),
+			Timestamp:  uint64(time.Now().Unix()),
+			Ttl:        uint64(10), // 10ms ttl
+		}
+	}
+
+	_, err = mgr.HandleHeartbeat(newHeartbeatReq(executorID1))
+	require.NoError(t, err)
+	_, err = mgr.HandleHeartbeat(newHeartbeatReq(executorID2))
+	require.NoError(t, err)
+
+	require.Equal(t, 2, mgr.ExecutorCount(model.Running))
+
+	require.Eventually(t, func() bool {
+		resp, err := mgr.HandleHeartbeat(newHeartbeatReq(executorID2))
+		require.NoError(t, err)
+		require.Nil(t, resp.Err)
+		return mgr.ExecutorCount(model.Running) == 1
+	}, time.Second*2, time.Millisecond*5)
+
+	event = <-stream.C
+	require.Equal(t, model.ExecutorStatusChange{
+		ID: executorID1,
+		Tp: model.EventExecutorOffline,
+	}, event)
+}