Merge pull request #30 from alibaba/bugfix-panic-mapreduce

Bugfix panic mapreduce

README.md

@@ -350,6 +350,11 @@ func (mr *TestMapReduceJob) Reduce(jobCtx *jobcontext.JobContext) (*processor.Pr
 package main
 
 import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"github.com/alibaba/schedulerx-worker-go"
 )
 
@@ -379,7 +384,12 @@ func main() {
 	client.RegisterTask("TestBroadcast", task2)
 	client.RegisterTask("TestMapJob", task3)
 	client.RegisterTask("TestMapReduceJob", task4)
-	select {}
+	
+	// wait for the stop signal
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
+	<-c
+	time.Sleep(time.Second * 5)
 }
 
 ```

client.go

@@ -21,11 +21,12 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 
-	"github.com/asynkron/protoactor-go/actor"
-
 	"github.com/alibaba/schedulerx-worker-go/config"
 	sxactor "github.com/alibaba/schedulerx-worker-go/internal/actor"
 	actorcomm "github.com/alibaba/schedulerx-worker-go/internal/actor/common"
@@ -47,12 +48,21 @@ var (
 )
 
 type Client struct {
-	cfg            *Config
-	opts           *Options
-	connpool       pool.ConnPool
-	tasks          *tasks.TaskMap
-	actorSystem    *actor.ActorSystem
-	taskMasterPool *masterpool.TaskMasterPool
+	cfg   *Config
+	opts  *Options
+	tasks *tasks.TaskMap
+
+	stopChan chan os.Signal
+}
+
+func (c *Client) RegisterTask(name string, task processor.Processor) {
+	c.tasks.Register(name, task)
+}
+
+// Shutdown gracefully stop the client
+func (c *Client) Shutdown() {
+	c.stopChan <- syscall.SIGINT
+	time.Sleep(time.Second * 5)
 }
 
 type Config struct {
@@ -149,23 +159,22 @@ func newClient(cfg *Config, opts ...Option) (*Client, error) {
 
 	// Init actors
 	actorSystem := actorcomm.GetActorSystem()
-	if err := sxactor.InitActors(actorSystem); err != nil {
-		return nil, fmt.Errorf("Init actors faild, err=%s. ", err.Error())
+	if err = sxactor.InitActors(actorSystem); err != nil {
+		return nil, fmt.Errorf("init actors faild, err=%s", err.Error())
 	}
 
+	stopChan := make(chan os.Signal, 1)
+	signal.Notify(stopChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
+
 	// Keep heartbeat, and receive message
 	// KeepHeartbeat must after init actors, so that can get actorSystemPort from actorSystem
-	go remoting.KeepHeartbeat(ctx, actorSystem, cfg.AppKey)
+	go remoting.KeepHeartbeat(ctx, actorSystem, cfg.AppKey, stopChan)
 	go remoting.OnMsgReceived(ctx)
 
 	return &Client{
-		cfg:         cfg,
-		opts:        options,
-		tasks:       taskMap,
-		actorSystem: actorSystem,
+		cfg:      cfg,
+		opts:     options,
+		tasks:    taskMap,
+		stopChan: stopChan,
 	}, nil
 }
-
-func (c *Client) RegisterTask(name string, task processor.Processor) {
-	c.tasks.Register(name, task)
-}

example/broadcast/main.go

@@ -17,6 +17,11 @@
 package main
 
 import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"github.com/alibaba/schedulerx-worker-go"
 )
 
@@ -36,5 +41,10 @@ func main() {
 	// The name TestBroadcast registered here must be consistent with the configured on the platform
 	task := &TestBroadcast{}
 	client.RegisterTask("TestBroadcast", task)
-	select {}
+
+	// wait for the stop signal
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
+	<-c
+	time.Sleep(time.Second * 5)
 }

example/broadcast/test_broadcast.go

@@ -31,7 +31,7 @@ var _ processor.BroadcastProcessor = &TestBroadcast{}
 type TestBroadcast struct{}
 
 // Process all machines would execute it.
-func (t TestBroadcast) Process(ctx *jobcontext.JobContext) (*processor.ProcessResult, error) {
+func (t *TestBroadcast) Process(ctx *jobcontext.JobContext) (*processor.ProcessResult, error) {
 	value := rand.Intn(10)
 	fmt.Printf("Total sharding num=%d, sharding=%d, value=%d\n", ctx.ShardingNum(), ctx.ShardingId(), value)
 	ret := new(processor.ProcessResult)
@@ -41,13 +41,13 @@ func (t TestBroadcast) Process(ctx *jobcontext.JobContext) (*processor.ProcessRe
 }
 
 // PreProcess only one machine will execute it.
-func (t TestBroadcast) PreProcess(ctx *jobcontext.JobContext) error {
+func (t *TestBroadcast) PreProcess(ctx *jobcontext.JobContext) error {
 	fmt.Println("TestBroadcastJob PreProcess")
 	return nil
 }
 
 // PostProcess only one machine will execute it.
-func (t TestBroadcast) PostProcess(ctx *jobcontext.JobContext) (*processor.ProcessResult, error) {
+func (t *TestBroadcast) PostProcess(ctx *jobcontext.JobContext) (*processor.ProcessResult, error) {
 	fmt.Println("TestBroadcastJob PostProcess")
 	allTaskResults := ctx.TaskResults()
 	allTaskStatuses := ctx.TaskStatuses()
@@ -68,7 +68,7 @@ func (t TestBroadcast) PostProcess(ctx *jobcontext.JobContext) (*processor.Proce
 	return ret, nil
 }
 
-func (t TestBroadcast) Kill(ctx *jobcontext.JobContext) error {
+func (t *TestBroadcast) Kill(ctx *jobcontext.JobContext) error {
 	fmt.Println("[Kill] Start kill my task: TestBroadcast")
 	return nil
 }

example/mapreduce/main.go

@@ -17,6 +17,11 @@
 package main
 
 import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"github.com/alibaba/schedulerx-worker-go"
 	"github.com/alibaba/schedulerx-worker-go/processor/mapjob"
 )
@@ -39,5 +44,10 @@ func main() {
 		mapjob.NewMapReduceJobProcessor(), // FIXME how define user behavior
 	}
 	client.RegisterTask("TestMapReduceJob", task)
-	select {}
+
+	// wait for the stop signal
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
+	<-c
+	time.Sleep(time.Second * 5)
 }

example/standalone/main.go

@@ -17,6 +17,11 @@
 package main
 
 import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"github.com/alibaba/schedulerx-worker-go"
 )
 
@@ -38,5 +43,10 @@ func main() {
 	// The name HelloWorld registered here must be consistent with the configured on the platform
 	task := &HelloWorld{}
 	client.RegisterTask("HelloWorld", task)
-	select {}
+
+	// wait for the stop signal
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
+	<-c
+	time.Sleep(time.Second * 5)
 }

internal/batch/base_req_handler.go

@@ -177,7 +177,6 @@ func (rcvr *BaseReqHandler) Stop() {
 	}
 	if rcvr.reqsQueue != nil {
 		rcvr.reqsQueue.Clear()
-		rcvr.reqsQueue = nil
 	}
 	if rcvr.activeRunnableNum != nil {
 		rcvr.activeRunnableNum = nil

internal/remoting/heartbeat.go

@@ -23,14 +23,12 @@ import (
 	"fmt"
 	"math"
 	"os"
-	"os/signal"
 	"runtime"
 	"syscall"
 	"time"
 
 	"github.com/asynkron/protoactor-go/actor"
 	"github.com/shirou/gopsutil/load"
-	"go.uber.org/atomic"
 	"google.golang.org/protobuf/proto"
 
 	"github.com/alibaba/schedulerx-worker-go/config"
@@ -51,22 +49,21 @@ var (
 	waitHeartbeatRespTimeout = 5 * time.Second
 )
 
-func KeepHeartbeat(ctx context.Context, actorSystem *actor.ActorSystem, appKey string) {
+func KeepHeartbeat(ctx context.Context, actorSystem *actor.ActorSystem, appKey string, stopChan chan os.Signal) {
 	var (
-		online         = atomic.NewBool(true)
 		taskMasterPool = masterpool.GetTaskMasterPool()
 		groupManager   = discovery.GetGroupManager()
 	)
 
-	heartbeat := func() {
+	heartbeat := func(online bool) {
 		_, actorSystemPort, err := actorSystem.GetHostPort()
 		if err != nil {
 			logger.Errorf("Write heartbeat to remote failed due to get actorSystem port failed, err=%s", err.Error())
 			return
 		}
 		for groupId, appGroupId := range groupManager.GroupId2AppGroupIdMap() {
 			jobInstanceIds := taskMasterPool.GetInstanceIds(appGroupId)
-			heartbeatReq := genHeartBeatRequest(groupId, appGroupId, jobInstanceIds, actorSystemPort, online.Load(), appKey)
+			heartbeatReq := genHeartBeatRequest(groupId, appGroupId, jobInstanceIds, actorSystemPort, online, appKey)
 			if err := sendHeartbeat(ctx, heartbeatReq); err != nil {
 				if errors.Is(err, syscall.EPIPE) || errors.Is(err, os.ErrDeadlineExceeded) {
 					pool.GetConnPool().ReconnectTrigger() <- struct{}{}
@@ -77,22 +74,19 @@ func KeepHeartbeat(ctx context.Context, actorSystem *actor.ActorSystem, appKey s
 			logger.Debugf("Write heartbeat to remote succeed.")
 		}
 	}
-	heartbeat()
-
-	// send worker offline heartbeat when shutdown
-	go func() {
-		c := make(chan os.Signal, 1)
-		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)
-		<-c
-		online.Store(false)
-		heartbeat()
-		logger.Infof("Write shutdown heartbeat to remote succeed.")
-	}()
+	heartbeat(true)
 
 	ticker := time.NewTicker(heartbeatInterval)
 	defer ticker.Stop()
-	for range ticker.C {
-		heartbeat()
+	for {
+		select {
+		case <-stopChan:
+			heartbeat(false)
+			logger.Infof("Write shutdown heartbeat to remote succeed.")
+			return
+		case <-ticker.C:
+			heartbeat(true)
+		}
 	}
 }