github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/background/worker.go (about)

     1  package background
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"runtime/pprof"
     7  	"sync"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
    13  )
    14  
    15  var (
    16  	ErrAlreadyClosed       = xerrors.Wrap(errors.New("ydb: background worker already closed"))
    17  	errClosedWithNilReason = xerrors.Wrap(errors.New("ydb: background worker closed with nil reason"))
    18  )
    19  
    20  // A Worker must not be copied after first use
    21  type Worker struct {
    22  	ctx            context.Context
    23  	workers        sync.WaitGroup
    24  	closeReason    error
    25  	tasksCompleted empty.Chan
    26  	tasks          chan backgroundTask
    27  	stop           context.CancelFunc
    28  	onceInit       sync.Once
    29  	m              xsync.Mutex
    30  	closed         bool
    31  }
    32  
    33  type CallbackFunc func(ctx context.Context)
    34  
    35  func NewWorker(parent context.Context) *Worker {
    36  	w := Worker{}
    37  	w.ctx, w.stop = xcontext.WithCancel(parent)
    38  
    39  	return &w
    40  }
    41  
    42  func (b *Worker) Context() context.Context {
    43  	b.init()
    44  
    45  	return b.ctx
    46  }
    47  
    48  func (b *Worker) Start(name string, f CallbackFunc) {
    49  	b.init()
    50  
    51  	b.m.WithLock(func() {
    52  		if b.closed {
    53  			return
    54  		}
    55  
    56  		b.tasks <- backgroundTask{
    57  			callback: f,
    58  			name:     name,
    59  		}
    60  	})
    61  }
    62  
    63  func (b *Worker) Done() <-chan struct{} {
    64  	b.init()
    65  
    66  	return b.ctx.Done()
    67  }
    68  
    69  func (b *Worker) Close(ctx context.Context, err error) error {
    70  	b.init()
    71  
    72  	var resErr error
    73  	b.m.WithLock(func() {
    74  		if b.closed {
    75  			resErr = xerrors.WithStackTrace(ErrAlreadyClosed)
    76  
    77  			return
    78  		}
    79  
    80  		b.closed = true
    81  
    82  		close(b.tasks)
    83  		b.closeReason = err
    84  		if b.closeReason == nil {
    85  			b.closeReason = errClosedWithNilReason
    86  		}
    87  
    88  		b.stop()
    89  	})
    90  	if resErr != nil {
    91  		return resErr
    92  	}
    93  
    94  	<-b.tasksCompleted
    95  
    96  	bgCompleted := make(empty.Chan)
    97  
    98  	go func() {
    99  		b.workers.Wait()
   100  		close(bgCompleted)
   101  	}()
   102  
   103  	select {
   104  	case <-bgCompleted:
   105  		return nil
   106  	case <-ctx.Done():
   107  		return ctx.Err()
   108  	}
   109  }
   110  
   111  func (b *Worker) CloseReason() error {
   112  	b.m.Lock()
   113  	defer b.m.Unlock()
   114  
   115  	return b.closeReason
   116  }
   117  
   118  func (b *Worker) init() {
   119  	b.onceInit.Do(func() {
   120  		if b.ctx == nil {
   121  			b.ctx, b.stop = xcontext.WithCancel(context.Background())
   122  		}
   123  		b.tasks = make(chan backgroundTask)
   124  		b.tasksCompleted = make(empty.Chan)
   125  		go b.starterLoop()
   126  	})
   127  }
   128  
   129  func (b *Worker) starterLoop() {
   130  	defer close(b.tasksCompleted)
   131  
   132  	for bgTask := range b.tasks {
   133  		b.workers.Add(1)
   134  
   135  		go func(task backgroundTask) {
   136  			defer b.workers.Done()
   137  
   138  			pprof.Do(b.ctx, pprof.Labels("background", task.name), task.callback)
   139  		}(bgTask)
   140  	}
   141  }
   142  
   143  type backgroundTask struct {
   144  	callback CallbackFunc
   145  	name     string
   146  }