github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/jobmaster/example/worker_impl.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package example
    15  
    16  import (
    17  	"context"
    18  	"strconv"
    19  	"sync"
    20  	"time"
    21  
    22  	"github.com/pingcap/log"
    23  	"github.com/pingcap/tiflow/engine/framework"
    24  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    25  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    26  	"go.uber.org/zap"
    27  )
    28  
    29  var _ framework.Worker = &exampleWorker{}
    30  
    31  var (
    32  	tickKey   = "tick_count"
    33  	testTopic = "test_topic"
    34  	testMsg   = "test_msg"
    35  )
    36  
    37  type exampleWorker struct {
    38  	framework.BaseWorker
    39  
    40  	work struct {
    41  		mu        sync.Mutex
    42  		tickCount int
    43  		finished  bool
    44  	}
    45  	wg sync.WaitGroup
    46  }
    47  
    48  func (w *exampleWorker) run() {
    49  	defer w.wg.Done()
    50  
    51  	time.Sleep(time.Second)
    52  	w.work.mu.Lock()
    53  	count := w.work.tickCount
    54  	w.work.mu.Unlock()
    55  	// nolint:errcheck
    56  	_, _ = w.BaseWorker.MetaKVClient().Put(
    57  		context.TODO(), tickKey, strconv.Itoa(count))
    58  
    59  	w.work.mu.Lock()
    60  	w.work.finished = true
    61  	w.work.mu.Unlock()
    62  
    63  	// nolint:errcheck
    64  	_ = w.BaseWorker.SendMessage(context.TODO(), testTopic, testMsg, true)
    65  }
    66  
    67  func (w *exampleWorker) InitImpl(ctx context.Context) error {
    68  	log.Info("InitImpl")
    69  	w.wg.Add(1)
    70  	go w.run()
    71  	return nil
    72  }
    73  
    74  func (w *exampleWorker) Tick(ctx context.Context) error {
    75  	log.Info("Tick")
    76  	w.work.mu.Lock()
    77  	w.work.tickCount++
    78  	count := w.work.tickCount
    79  	w.work.mu.Unlock()
    80  
    81  	storage, err := w.OpenStorage(ctx, "/local/example")
    82  	if err != nil {
    83  		return err
    84  	}
    85  
    86  	file, err := storage.BrExternalStorage().Create(ctx, strconv.Itoa(count)+".txt", nil)
    87  	if err != nil {
    88  		return err
    89  	}
    90  	_, err = file.Write(ctx, []byte(strconv.Itoa(count)))
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	if err := file.Close(ctx); err != nil {
    96  		return err
    97  	}
    98  	return storage.Persist(ctx)
    99  }
   100  
   101  func (w *exampleWorker) Status() frameModel.WorkerStatus {
   102  	log.Info("Status")
   103  	code := frameModel.WorkerStateNormal
   104  	w.work.mu.Lock()
   105  	finished := w.work.finished
   106  	w.work.mu.Unlock()
   107  
   108  	if finished {
   109  		code = frameModel.WorkerStateFinished
   110  	}
   111  	return frameModel.WorkerStatus{State: code}
   112  }
   113  
   114  func (w *exampleWorker) OnMasterMessage(ctx context.Context, topic p2p.Topic, message p2p.MessageValue) error {
   115  	log.Info("OnMasterMessage", zap.Any("message", message))
   116  	return nil
   117  }
   118  
   119  func (w *exampleWorker) CloseImpl(ctx context.Context) {
   120  	log.Info("CloseImpl")
   121  	w.wg.Wait()
   122  }