124 lines
3.2 KiB
Go
124 lines
3.2 KiB
Go
package postgres
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/metric/instrument/syncfloat64"
|
|
"log"
|
|
"openreplay/backend/pkg/monitoring"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
insertPrefix = `INSERT INTO `
|
|
insertValues = ` VALUES `
|
|
insertSuffix = ` ON CONFLICT DO NOTHING;`
|
|
)
|
|
|
|
type Bulk interface {
|
|
Append(args ...interface{}) error
|
|
Send() error
|
|
Table() string
|
|
}
|
|
|
|
type bulkImpl struct {
|
|
conn Pool
|
|
table string
|
|
columns string
|
|
template string
|
|
setSize int
|
|
sizeLimit int
|
|
values []interface{}
|
|
bulkSize syncfloat64.Histogram
|
|
bulkDuration syncfloat64.Histogram
|
|
}
|
|
|
|
func (b *bulkImpl) Append(args ...interface{}) error {
|
|
if len(args) != b.setSize {
|
|
return fmt.Errorf("wrong number of arguments, waited: %d, got: %d", b.setSize, len(args))
|
|
}
|
|
b.values = append(b.values, args...)
|
|
if len(b.values)/b.setSize >= b.sizeLimit {
|
|
return b.send()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (b *bulkImpl) Send() error {
|
|
if len(b.values) == 0 {
|
|
return nil
|
|
}
|
|
return b.send()
|
|
}
|
|
|
|
func (b *bulkImpl) Table() string {
|
|
return b.table
|
|
}
|
|
|
|
func (b *bulkImpl) send() error {
|
|
start := time.Now()
|
|
size := len(b.values) / b.setSize
|
|
request := bytes.NewBufferString(insertPrefix + b.table + b.columns + insertValues)
|
|
args := make([]interface{}, b.setSize)
|
|
for i := 0; i < len(b.values)/b.setSize; i++ {
|
|
for j := 0; j < b.setSize; j++ {
|
|
args[j] = i*b.setSize + j + 1
|
|
}
|
|
if i > 0 {
|
|
request.WriteByte(',')
|
|
}
|
|
request.WriteString(fmt.Sprintf(b.template, args...))
|
|
}
|
|
request.WriteString(insertSuffix)
|
|
err := b.conn.Exec(request.String(), b.values...)
|
|
b.values = make([]interface{}, 0, b.setSize*b.sizeLimit)
|
|
if err != nil {
|
|
return fmt.Errorf("send bulk err: %s", err)
|
|
}
|
|
// Save bulk metrics
|
|
ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*200)
|
|
b.bulkDuration.Record(ctx, float64(time.Now().Sub(start).Milliseconds()), attribute.String("table", b.table))
|
|
b.bulkSize.Record(ctx, float64(size), attribute.String("table", b.table))
|
|
return nil
|
|
}
|
|
|
|
func NewBulk(conn Pool, metrics *monitoring.Metrics, table, columns, template string, setSize, sizeLimit int) (Bulk, error) {
|
|
switch {
|
|
case conn == nil:
|
|
return nil, errors.New("db conn is empty")
|
|
case metrics == nil:
|
|
return nil, errors.New("metrics is empty")
|
|
case table == "":
|
|
return nil, errors.New("table is empty")
|
|
case columns == "":
|
|
return nil, errors.New("columns is empty")
|
|
case template == "":
|
|
return nil, errors.New("template is empty")
|
|
case setSize <= 0:
|
|
return nil, errors.New("set size is wrong")
|
|
case sizeLimit <= 0:
|
|
return nil, errors.New("size limit is wrong")
|
|
}
|
|
messagesInBulk, err := metrics.RegisterHistogram("messages_in_bulk")
|
|
if err != nil {
|
|
log.Printf("can't create messages_size metric: %s", err)
|
|
}
|
|
bulkInsertDuration, err := metrics.RegisterHistogram("bulk_insert_duration")
|
|
if err != nil {
|
|
log.Printf("can't create messages_size metric: %s", err)
|
|
}
|
|
return &bulkImpl{
|
|
conn: conn,
|
|
table: table,
|
|
columns: columns,
|
|
template: template,
|
|
setSize: setSize,
|
|
sizeLimit: sizeLimit,
|
|
values: make([]interface{}, 0, setSize*sizeLimit),
|
|
bulkSize: messagesInBulk,
|
|
bulkDuration: bulkInsertDuration,
|
|
}, nil
|
|
}
|