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

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