github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/pkg/parallel-writer/writer.go (about)

     1  // Copyright (c) 2021 Terminus, 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 writer
    16  
    17  import (
    18  	"time"
    19  
    20  	"github.com/recallsong/go-utils/errorx"
    21  )
    22  
    23  // Writer .
    24  type Writer interface {
    25  	Write(data interface{}) error
    26  	WriteN(data ...interface{}) (int, error)
    27  	Close() error
    28  }
    29  
    30  // ErrorHandler .
    31  type ErrorHandler func(error) error
    32  
    33  // IngoreError .
    34  func IngoreError(error) error { return nil }
    35  
    36  // ErrorAbort .
    37  func ErrorAbort(err error) error { return err }
    38  
    39  // ParallelBatch .
    40  func ParallelBatch(
    41  	writers func(i uint64) Writer, // get writer function for each goroutine
    42  	parallelism, // goroutine number
    43  	size uint64, // buffer size
    44  	timeout time.Duration, // timeout for buffer flush
    45  	errorh ErrorHandler, // error handler
    46  ) Writer {
    47  	if parallelism <= 0 {
    48  		if size <= 1 {
    49  			return writers(0)
    50  		}
    51  		parallelism = 1
    52  	}
    53  	if size <= 1 {
    54  		writer := &channelWriter{
    55  			dataCh:      make(chan interface{}, parallelism),
    56  			errorCh:     make(chan error, parallelism),
    57  			parallelism: parallelism,
    58  		}
    59  		for i := uint64(0); i < parallelism; i++ {
    60  			go func(w Writer, in *channelWriter) {
    61  				var err error
    62  				defer func() {
    63  					cerr := w.Close()
    64  					if err == nil {
    65  						err = cerr
    66  					}
    67  					writer.errorCh <- err
    68  				}()
    69  				for data := range in.dataCh {
    70  					err = w.Write(data)
    71  					if err != nil {
    72  						if errorh != nil {
    73  							err = errorh(err)
    74  							if err != nil {
    75  								return
    76  							}
    77  						} else {
    78  							return
    79  						}
    80  					}
    81  				}
    82  			}(writers(i), writer)
    83  		}
    84  		return writer
    85  	}
    86  	writer := &channelWriter{
    87  		dataCh:      make(chan interface{}, parallelism*size),
    88  		errorCh:     make(chan error, parallelism),
    89  		parallelism: parallelism,
    90  	}
    91  	for i := uint64(0); i < parallelism; i++ {
    92  		go func(w Writer, in *channelWriter) {
    93  			buf := NewBuffer(w, int(size))
    94  			tick := time.NewTicker(timeout)
    95  			var err error
    96  			defer func() {
    97  				tick.Stop()
    98  				cerr := buf.Close()
    99  				if cerr != nil {
   100  					cerr = errorh(cerr)
   101  					if err == nil {
   102  						err = cerr
   103  					}
   104  				}
   105  				writer.errorCh <- err
   106  			}()
   107  			for {
   108  				select {
   109  				case data, ok := <-in.dataCh:
   110  					if !ok {
   111  						return
   112  					}
   113  					err = buf.Write(data)
   114  					if err != nil {
   115  						if errorh != nil {
   116  							err = errorh(err)
   117  							if err != nil {
   118  								return
   119  							}
   120  						} else {
   121  							return
   122  						}
   123  					}
   124  				case <-tick.C:
   125  					err = buf.Flush()
   126  					if err != nil {
   127  						if errorh != nil {
   128  							err = errorh(err)
   129  							if err != nil {
   130  								return
   131  							}
   132  						} else {
   133  							return
   134  						}
   135  					}
   136  				}
   137  			}
   138  		}(writers(i), writer)
   139  	}
   140  	return writer
   141  }
   142  
   143  type channelWriter struct {
   144  	dataCh      chan interface{}
   145  	errorCh     chan error
   146  	parallelism uint64
   147  }
   148  
   149  func (w *channelWriter) Write(data interface{}) error {
   150  	w.dataCh <- data
   151  	return nil
   152  }
   153  
   154  func (w *channelWriter) WriteN(data ...interface{}) (int, error) {
   155  	for _, item := range data {
   156  		w.dataCh <- item
   157  	}
   158  	return len(data), nil
   159  }
   160  
   161  func (w *channelWriter) Close() error {
   162  	close(w.dataCh)
   163  	var errs errorx.Errors
   164  	for i := uint64(0); i < w.parallelism; i++ {
   165  		err := <-w.errorCh
   166  		if err != nil {
   167  			errs = append(errs, err)
   168  		}
   169  	}
   170  	close(w.errorCh)
   171  	return errs.MaybeUnwrap()
   172  }