github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/dbs/batch.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package dbs
     4  
     5  import (
     6  	"database/sql"
     7  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
     8  	"time"
     9  )
    10  
    11  type batchItem struct {
    12  	query string
    13  	args  []any
    14  }
    15  
    16  type Batch struct {
    17  	db *DB
    18  	n  int
    19  
    20  	enableStat bool
    21  
    22  	onFail func(err error)
    23  
    24  	queue      chan *batchItem
    25  	closeEvent chan bool
    26  
    27  	isClosed bool
    28  }
    29  
    30  func NewBatch(db *DB, n int) *Batch {
    31  	var batch = &Batch{
    32  		db:         db,
    33  		n:          n,
    34  		queue:      make(chan *batchItem, 16),
    35  		closeEvent: make(chan bool, 1),
    36  	}
    37  	db.batches = append(db.batches, batch)
    38  	return batch
    39  }
    40  
    41  func (this *Batch) EnableStat(b bool) {
    42  	this.enableStat = b
    43  }
    44  
    45  func (this *Batch) OnFail(callback func(err error)) {
    46  	this.onFail = callback
    47  }
    48  
    49  func (this *Batch) Add(query string, args ...any) {
    50  	if this.isClosed {
    51  		return
    52  	}
    53  	this.queue <- &batchItem{
    54  		query: query,
    55  		args:  args,
    56  	}
    57  }
    58  
    59  func (this *Batch) Exec() {
    60  	var n = this.n
    61  	if n <= 0 {
    62  		n = 4
    63  	}
    64  
    65  	var ticker = time.NewTicker(100 * time.Millisecond)
    66  	var count = 0
    67  	var lastTx *sql.Tx
    68  For:
    69  	for {
    70  		// closed
    71  		if this.isClosed {
    72  			if lastTx != nil {
    73  				_ = this.commitTx(lastTx)
    74  				lastTx = nil
    75  			}
    76  
    77  			return
    78  		}
    79  
    80  		select {
    81  		case item := <-this.queue:
    82  			if lastTx == nil {
    83  				lastTx = this.beginTx()
    84  				if lastTx == nil {
    85  					continue For
    86  				}
    87  			}
    88  
    89  			err := this.execItem(lastTx, item)
    90  			if err != nil {
    91  				if IsClosedErr(err) {
    92  					return
    93  				}
    94  				this.processErr(item.query, err)
    95  			}
    96  
    97  			count++
    98  
    99  			if count == n {
   100  				count = 0
   101  				err = this.commitTx(lastTx)
   102  				lastTx = nil
   103  				if err != nil {
   104  					if IsClosedErr(err) {
   105  						return
   106  					}
   107  					this.processErr("commit", err)
   108  				}
   109  			}
   110  		case <-ticker.C:
   111  			if lastTx == nil || count == 0 {
   112  				continue For
   113  			}
   114  			count = 0
   115  			err := this.commitTx(lastTx)
   116  			lastTx = nil
   117  			if err != nil {
   118  				if IsClosedErr(err) {
   119  					return
   120  				}
   121  				this.processErr("commit", err)
   122  			}
   123  		case <-this.closeEvent:
   124  			// closed
   125  
   126  			if lastTx != nil {
   127  				_ = this.commitTx(lastTx)
   128  				lastTx = nil
   129  			}
   130  
   131  			return
   132  		}
   133  	}
   134  }
   135  
   136  func (this *Batch) close() {
   137  	this.isClosed = true
   138  
   139  	select {
   140  	case this.closeEvent <- true:
   141  	default:
   142  
   143  	}
   144  }
   145  
   146  func (this *Batch) beginTx() *sql.Tx {
   147  	if !this.db.BeginUpdating() {
   148  		return nil
   149  	}
   150  
   151  	tx, err := this.db.Begin()
   152  	if err != nil {
   153  		this.processErr("begin transaction", err)
   154  		this.db.EndUpdating()
   155  		return nil
   156  	}
   157  	return tx
   158  }
   159  
   160  func (this *Batch) commitTx(tx *sql.Tx) error {
   161  	// always commit without checking database closing status
   162  	this.db.EndUpdating()
   163  	return tx.Commit()
   164  }
   165  
   166  func (this *Batch) execItem(tx *sql.Tx, item *batchItem) error {
   167  	// check database status
   168  	if this.db.BeginUpdating() {
   169  		defer this.db.EndUpdating()
   170  	} else {
   171  		return errDBIsClosed
   172  	}
   173  
   174  	if this.enableStat {
   175  		defer SharedQueryStatManager.AddQuery(item.query).End()
   176  	}
   177  
   178  	_, err := tx.Exec(item.query, item.args...)
   179  	return err
   180  }
   181  
   182  func (this *Batch) processErr(prefix string, err error) {
   183  	if err == nil {
   184  		return
   185  	}
   186  
   187  	if this.onFail != nil {
   188  		this.onFail(err)
   189  	} else {
   190  		remotelogs.Error("SQLITE_BATCH", prefix+": "+err.Error())
   191  	}
   192  }