github.com/dolthub/go-mysql-server@v0.18.0/sql/background_threads.go (about)

     1  // Copyright 2021 Dolthub, 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  // 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 sql
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"sync"
    21  )
    22  
    23  var ErrCannotAddToClosedBackgroundThreads = errors.New("cannot add to a close background threads instance")
    24  
    25  type BackgroundThreads struct {
    26  	wg           *sync.WaitGroup
    27  	mu           *sync.Mutex
    28  	parentCtx    context.Context
    29  	parentCancel context.CancelFunc
    30  	nameToCancel map[string]context.CancelFunc
    31  	nameToCtx    map[string]context.Context
    32  }
    33  
    34  func NewBackgroundThreads() *BackgroundThreads {
    35  	ctx, cancel := context.WithCancel(context.Background())
    36  	return &BackgroundThreads{
    37  		wg:           &sync.WaitGroup{},
    38  		parentCtx:    ctx,
    39  		parentCancel: cancel,
    40  		mu:           &sync.Mutex{},
    41  		nameToCancel: make(map[string]context.CancelFunc),
    42  		nameToCtx:    make(map[string]context.Context),
    43  	}
    44  }
    45  
    46  // Add starts a background goroutine wrapped by a top-level sync.WaitGroup.
    47  // [f] must return when its [ctx] argument is cancelled, otherwise
    48  // Shutdown will hang.
    49  func (bt *BackgroundThreads) Add(name string, f func(ctx context.Context)) error {
    50  	select {
    51  	case <-bt.parentCtx.Done():
    52  		return ErrCannotAddToClosedBackgroundThreads
    53  	default:
    54  	}
    55  
    56  	threadCtx, threadCancel := context.WithCancel(bt.parentCtx)
    57  
    58  	bt.mu.Lock()
    59  	defer bt.mu.Unlock()
    60  
    61  	bt.nameToCancel[name] = threadCancel
    62  	bt.nameToCtx[name] = threadCtx
    63  	bt.wg.Add(1)
    64  
    65  	go func() {
    66  		defer bt.wg.Done()
    67  		f(threadCtx)
    68  	}()
    69  
    70  	return nil
    71  }
    72  
    73  // Shutdown cancels the parent context for every async thread,
    74  // and waits for each goroutine to drain and return before exiting.
    75  func (bt *BackgroundThreads) Shutdown() error {
    76  	bt.parentCancel()
    77  	bt.wg.Wait()
    78  	return bt.parentCtx.Err()
    79  }