github.com/matrixorigin/matrixone@v0.7.0/pkg/vm/engine/tae/tasks/job.go (about)

     1  // Copyright 2021 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package tasks
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"time"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/logutil"
    23  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/common"
    24  	"github.com/panjf2000/ants/v2"
    25  )
    26  
    27  type JobScheduler interface {
    28  	Schedule(job *Job) error
    29  	Stop()
    30  }
    31  
    32  type JobExecutor = func(context.Context) *JobResult
    33  
    34  type JobResult struct {
    35  	Err error
    36  	Res any
    37  }
    38  
    39  var SerialJobScheduler = new(simpleJobSceduler)
    40  
    41  type simpleJobSceduler struct{}
    42  
    43  func (s *simpleJobSceduler) Stop() {}
    44  func (s *simpleJobSceduler) Schedule(job *Job) (err error) {
    45  	job.Run()
    46  	return
    47  }
    48  
    49  type parallelJobScheduler struct {
    50  	pool *ants.Pool
    51  }
    52  
    53  func NewParallelJobScheduler(parallism int) *parallelJobScheduler {
    54  	pool, err := ants.NewPool(parallism)
    55  	if err != nil {
    56  		panic(err)
    57  	}
    58  	return &parallelJobScheduler{
    59  		pool: pool,
    60  	}
    61  }
    62  
    63  func (s *parallelJobScheduler) Stop() {
    64  	s.pool.Release()
    65  	s.pool = nil
    66  }
    67  
    68  func (s *parallelJobScheduler) Schedule(job *Job) (err error) {
    69  	err = s.pool.Submit(job.Run)
    70  	return
    71  }
    72  
    73  type Job struct {
    74  	id      string
    75  	wg      *sync.WaitGroup
    76  	ctx     context.Context
    77  	exec    JobExecutor
    78  	result  *JobResult
    79  	startTs time.Time
    80  	endTs   time.Time
    81  }
    82  
    83  func NewJob(id string, ctx context.Context, exec JobExecutor) *Job {
    84  	e := &Job{
    85  		id:   id,
    86  		ctx:  ctx,
    87  		exec: exec,
    88  		wg:   new(sync.WaitGroup),
    89  	}
    90  	e.wg.Add(1)
    91  	return e
    92  }
    93  
    94  func (job *Job) Run() {
    95  	defer job.wg.Done()
    96  	job.startTs = time.Now()
    97  	defer func() {
    98  		job.endTs = time.Now()
    99  		logutil.Info("run-job", common.AnyField("id", job.id),
   100  			common.ErrorField(job.result.Err),
   101  			common.DurationField(job.endTs.Sub(job.startTs)))
   102  	}()
   103  	result := job.exec(job.ctx)
   104  	job.result = result
   105  }
   106  
   107  func (job *Job) WaitDone() *JobResult {
   108  	job.wg.Wait()
   109  	return job.result
   110  }
   111  
   112  func (job *Job) GetResult() *JobResult {
   113  	job.wg.Wait()
   114  	return job.result
   115  }
   116  
   117  func (job *Job) Close() {
   118  	job.result = nil
   119  }