cloudpods/pkg/cloudcommon/workmanager/manager.go

// Copyright 2019 Yunion
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package workmanager

import (
	"context"
	"runtime/debug"
	"sync/atomic"
	"time"

	"yunion.io/x/jsonutils"
	"yunion.io/x/log"

	"yunion.io/x/onecloud/pkg/appctx"
	"yunion.io/x/onecloud/pkg/appsrv"
)

type DelayTaskFunc func(context.Context, interface{}) (jsonutils.JSONObject, error)
type OnTaskFailed func(context.Context, string)
type OnTaskCompleted func(context.Context, jsonutils.JSONObject)

type SWorkManager struct {
	curCount int32

	onFailed    OnTaskFailed
	onCompleted OnTaskCompleted
	worker      *appsrv.SWorkerManager
}

func (w *SWorkManager) add() {
	atomic.AddInt32(&w.curCount, 1)
}

func (w *SWorkManager) done() {
	atomic.AddInt32(&w.curCount, -1)
}

func (w *SWorkManager) DelayTask(ctx context.Context, task DelayTaskFunc, params interface{}) {
	w.delayTask(ctx, task, params, w.worker)
}

func (w *SWorkManager) DelayTaskWithWorker(
	ctx context.Context, task DelayTaskFunc, params interface{}, worker *appsrv.SWorkerManager,
) {
	w.delayTask(ctx, task, params, worker)
}

// If delay task is not panic and task func return err is nil
// task complete will be called, otherwise called task failed
// Params is interface for receive any type, task func should do type assertion
func (w *SWorkManager) delayTask(ctx context.Context, task DelayTaskFunc, params interface{}, worker *appsrv.SWorkerManager) {
	if ctx == nil || ctx.Value(appctx.APP_CONTEXT_KEY_TASK_ID) == nil {
		w.delayTaskWithoutReqctx(ctx, task, params, worker)
		return
	} else {
		// delayTask should have a new context.Context with value 'taskid'
		ctx = context.WithValue(context.Background(), appctx.APP_CONTEXT_KEY_TASK_ID, ctx.Value(appctx.APP_CONTEXT_KEY_TASK_ID))
		w.add()
		worker.Run(func() {
			defer w.done()
			defer func() {
				if r := recover(); r != nil {
					log.Errorf("DelayTask panic: %s", r)
					debug.PrintStack()
					switch val := r.(type) {
					case string:
						w.onFailed(ctx, val)
					case error:
						w.onFailed(ctx, val.Error())
					default:
						w.onFailed(ctx, "Unknown panic")
					}
				}
			}()

			// HACK: callback only
			if task == nil {
				w.onCompleted(ctx, nil)
				return
			}

			res, err := task(ctx, params)
			if err != nil {
				log.Infof("DelayTask failed: %s", err)
				w.onFailed(ctx, err.Error())
			} else {
				log.Infof("DelayTask complete: %v", res)
				w.onCompleted(ctx, res)
			}
		}, nil, nil)
	}
}

func (w *SWorkManager) DelayTaskWithoutReqctx(ctx context.Context, task DelayTaskFunc, params interface{}) {
	w.delayTaskWithoutReqctx(ctx, task, params, w.worker)
}

// response task by self, did not callback
func (w *SWorkManager) delayTaskWithoutReqctx(
	ctx context.Context, task DelayTaskFunc, params interface{}, worker *appsrv.SWorkerManager,
) {
	// delayTaskWithoutReqctx should have a new context.Context
	newCtx := context.Background()
	if ctx != nil && ctx.Value(appctx.APP_CONTEXT_KEY_TASK_ID) != nil {
		newCtx = context.WithValue(newCtx, appctx.APP_CONTEXT_KEY_TASK_ID, ctx.Value(appctx.APP_CONTEXT_KEY_TASK_ID))
	}
	ctx = newCtx
	w.add()
	w.worker.Run(func() {
		defer w.done()
		defer func() {
			if r := recover(); r != nil {
				log.Errorln("DelayTaskWithoutReqctx panic: ", r)
				debug.PrintStack()
			}
		}()

		if task == nil {
			return
		}

		if _, err := task(ctx, params); err != nil {
			log.Errorln("DelayTaskWithoutReqctx error: ", err)
			w.onFailed(ctx, err.Error())
		}
	}, nil, nil)
}

func (w *SWorkManager) Stop() {
	log.Infof("WorkManager stop, waitting for workers ...")
	for w.curCount > 0 {
		log.Warningf("Busy workers count %d, waiting stopped", w.curCount)
		time.Sleep(1 * time.Second)
	}
}

func NewWorkManger(onFailed OnTaskFailed, onCompleted OnTaskCompleted, workerCount int) *SWorkManager {
	if workerCount <= 0 {
		workerCount = 1
	}
	return &SWorkManager{
		onFailed:    onFailed,
		onCompleted: onCompleted,
		worker: appsrv.NewWorkerManager(
			"RequestWorker", workerCount, appsrv.DEFAULT_BACKLOG, false),
	}
}