github.com/lingyao2333/mo-zero@v1.4.1/core/executors/chunkexecutor.go (about)

     1  package executors
     2  
     3  import "time"
     4  
     5  const defaultChunkSize = 1024 * 1024 // 1M
     6  
     7  type (
     8  	// ChunkOption defines the method to customize a ChunkExecutor.
     9  	ChunkOption func(options *chunkOptions)
    10  
    11  	// A ChunkExecutor is an executor to execute tasks when either requirement meets:
    12  	// 1. up to given chunk size
    13  	// 2. flush interval elapsed
    14  	ChunkExecutor struct {
    15  		executor  *PeriodicalExecutor
    16  		container *chunkContainer
    17  	}
    18  
    19  	chunkOptions struct {
    20  		chunkSize     int
    21  		flushInterval time.Duration
    22  	}
    23  )
    24  
    25  // NewChunkExecutor returns a ChunkExecutor.
    26  func NewChunkExecutor(execute Execute, opts ...ChunkOption) *ChunkExecutor {
    27  	options := newChunkOptions()
    28  	for _, opt := range opts {
    29  		opt(&options)
    30  	}
    31  
    32  	container := &chunkContainer{
    33  		execute:      execute,
    34  		maxChunkSize: options.chunkSize,
    35  	}
    36  	executor := &ChunkExecutor{
    37  		executor:  NewPeriodicalExecutor(options.flushInterval, container),
    38  		container: container,
    39  	}
    40  
    41  	return executor
    42  }
    43  
    44  // Add adds task with given chunk size into ce.
    45  func (ce *ChunkExecutor) Add(task interface{}, size int) error {
    46  	ce.executor.Add(chunk{
    47  		val:  task,
    48  		size: size,
    49  	})
    50  	return nil
    51  }
    52  
    53  // Flush forces ce to flush and execute tasks.
    54  func (ce *ChunkExecutor) Flush() {
    55  	ce.executor.Flush()
    56  }
    57  
    58  // Wait waits the execution to be done.
    59  func (ce *ChunkExecutor) Wait() {
    60  	ce.executor.Wait()
    61  }
    62  
    63  // WithChunkBytes customizes a ChunkExecutor with the given chunk size.
    64  func WithChunkBytes(size int) ChunkOption {
    65  	return func(options *chunkOptions) {
    66  		options.chunkSize = size
    67  	}
    68  }
    69  
    70  // WithFlushInterval customizes a ChunkExecutor with the given flush interval.
    71  func WithFlushInterval(duration time.Duration) ChunkOption {
    72  	return func(options *chunkOptions) {
    73  		options.flushInterval = duration
    74  	}
    75  }
    76  
    77  func newChunkOptions() chunkOptions {
    78  	return chunkOptions{
    79  		chunkSize:     defaultChunkSize,
    80  		flushInterval: defaultFlushInterval,
    81  	}
    82  }
    83  
    84  type chunkContainer struct {
    85  	tasks        []interface{}
    86  	execute      Execute
    87  	size         int
    88  	maxChunkSize int
    89  }
    90  
    91  func (bc *chunkContainer) AddTask(task interface{}) bool {
    92  	ck := task.(chunk)
    93  	bc.tasks = append(bc.tasks, ck.val)
    94  	bc.size += ck.size
    95  	return bc.size >= bc.maxChunkSize
    96  }
    97  
    98  func (bc *chunkContainer) Execute(tasks interface{}) {
    99  	vals := tasks.([]interface{})
   100  	bc.execute(vals)
   101  }
   102  
   103  func (bc *chunkContainer) RemoveAll() interface{} {
   104  	tasks := bc.tasks
   105  	bc.tasks = nil
   106  	bc.size = 0
   107  	return tasks
   108  }
   109  
   110  type chunk struct {
   111  	val  interface{}
   112  	size int
   113  }