github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/tasks/worker/heartbeater.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 ops
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"time"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/tasks/worker/base"
    23  )
    24  
    25  type lamdaHandle struct {
    26  	onExec func()
    27  	onStop func()
    28  }
    29  
    30  func (h *lamdaHandle) OnExec() {
    31  	if h.onExec != nil {
    32  		h.onExec()
    33  	}
    34  }
    35  
    36  func (h *lamdaHandle) OnStopped() {
    37  	if h.onStop != nil {
    38  		h.onStop()
    39  	}
    40  }
    41  
    42  type heartbeater struct {
    43  	handle   base.IHBHandle
    44  	interval time.Duration
    45  	ctx      context.Context
    46  	cancel   context.CancelFunc
    47  	wg       *sync.WaitGroup
    48  }
    49  
    50  func NewHeartBeaterWithFunc(interval time.Duration, onExec, onStop func()) *heartbeater {
    51  	h := &lamdaHandle{onExec: onExec, onStop: onStop}
    52  	return NewHeartBeater(interval, h)
    53  }
    54  
    55  func NewHeartBeater(interval time.Duration, handle base.IHBHandle) *heartbeater {
    56  	c := &heartbeater{
    57  		interval: interval,
    58  		handle:   handle,
    59  	}
    60  	c.ctx, c.cancel = context.WithCancel(context.Background())
    61  	return c
    62  }
    63  
    64  func (c *heartbeater) Start() {
    65  	c.wg = &sync.WaitGroup{}
    66  	c.wg.Add(1)
    67  	go func() {
    68  		defer c.wg.Done()
    69  		ticker := time.NewTicker(c.interval)
    70  		for {
    71  			select {
    72  			case <-c.ctx.Done():
    73  				ticker.Stop()
    74  				return
    75  			case <-ticker.C:
    76  				c.handle.OnExec()
    77  			}
    78  		}
    79  	}()
    80  }
    81  
    82  func (c *heartbeater) Stop() {
    83  	c.cancel()
    84  	c.wg.Wait()
    85  	c.handle.OnStopped()
    86  }