[IMPROVED] Bulk insert of messages

A new configuration `bulk_insert_limit` switches the server from
the current insertion of messages within a SQL transaction, to
an "INSERT INTO MESSAGES () VALUES (),(),..." which can speed up
performance by several folds.

The server still may perform regular insert within transaction
if the limit is deemed too low.

This new configuration parameter is not enabled by default. It
needs to be explicitly set, either in configuration file or
from command line `--sql_bulk_insert_limit <number here>`.

Resolves #1132

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>

nats-streaming-server.go

@@ -91,6 +91,7 @@ Streaming Server SQL Store Options:
     --sql_source <string>            Datasource used when opening an SQL connection to the database
     --sql_no_caching <bool>          Enable/Disable caching for improved performance
     --sql_max_open_conns <int>       Maximum number of opened connections to the database
+    --sql_bulk_insert_limit <int>    Maximum number of messages stored with a single SQL "INSERT" statement
 
 Streaming Server TLS Options:
     -secure <bool>                   Use a TLS connection to the NATS server without

server/conf.go

@@ -595,6 +595,11 @@ func parseSQLOptions(itf interface{}, opts *Options) error {
 				return err
 			}
 			opts.SQLStoreOpts.MaxOpenConns = int(v.(int64))
+		case "bulk_insert_limit":
+			if err := checkType(name, reflect.Int64, v); err != nil {
+				return err
+			}
+			opts.SQLStoreOpts.BulkInsertLimit = int(v.(int64))
 		}
 	}
 	return nil
@@ -688,6 +693,7 @@ func ConfigureOptions(fs *flag.FlagSet, args []string, printVersion, printHelp,
 	defSQLOpts := stores.DefaultSQLStoreOptions()
 	fs.BoolVar(&sopts.SQLStoreOpts.NoCaching, "sql_no_caching", defSQLOpts.NoCaching, "Enable/Disable caching")
 	fs.IntVar(&sopts.SQLStoreOpts.MaxOpenConns, "sql_max_open_conns", defSQLOpts.MaxOpenConns, "Max opened connections to the database")
+	fs.IntVar(&sopts.SQLStoreOpts.BulkInsertLimit, "sql_bulk_insert_limit", 0, "Limit the number of messages inserted in one SQL query")
 	fs.StringVar(&sopts.SyslogName, "syslog_name", "", "Syslog Name")
 	fs.BoolVar(&sopts.Encrypt, "encrypt", false, "Specify if server should use encryption at rest")
 	fs.StringVar(&sopts.EncryptionCipher, "encryption_cipher", stores.CryptoCipherAutoSelect, "Encryption cipher. Supported are AES and CHACHA (default is AES)")

server/conf_test.go

@@ -287,6 +287,9 @@ func TestParseConfig(t *testing.T) {
 	if opts.SQLStoreOpts.MaxOpenConns != 5 {
 		t.Fatalf("Expected SQL MaxOpenConns to be 5, got %v", opts.SQLStoreOpts.MaxOpenConns)
 	}
+	if opts.SQLStoreOpts.BulkInsertLimit != 1000 {
+		t.Fatalf("Expected SQL BulkInsertLimit to be 1000, got %v", opts.SQLStoreOpts.BulkInsertLimit)
+	}
 	if !opts.Encrypt {
 		t.Fatal("Expected Encrypt to be true")
 	}

stores/sqlstore.go

@@ -185,6 +185,10 @@ const (
 	// Limit of number of messages in the cache before message store
 	// is automatically flushed on a Store() call.
 	sqlDefaultMsgCacheLimit = 1024
+
+	// If bulk insert limit is set, the server will still insert messages
+	// using tx if the limit is below this threshold.
+	sqlMinBulkInsertLimit = 5
 )
 
 // These are initialized based on the constants that have reasonable values.
@@ -215,6 +219,11 @@ type SQLStoreOptions struct {
 	// APIs will cause execution of their respective SQL statements.
 	NoCaching bool
 
+	// If this is non 0, and NoCaching is not enabled, the server will perform
+	// bulk insert of messages. This is the limit of values added to the SQL statement
+	// "INSERT INTO Messages (..) VALUES (..)[,(..)*]".
+	BulkInsertLimit int
+
 	// Maximum number of open connections to the database.
 	// If <= 0, then there is no limit on the number of open connections.
 	// The default is 0 (unlimited).
@@ -240,6 +249,14 @@ func SQLNoCaching(noCaching bool) SQLStoreOption {
 	}
 }
 
+// SQLBulkInsertLimit sets the BulkInsertLimit option
+func SQLBulkInsertLimit(limit int) SQLStoreOption {
+	return func(o *SQLStoreOptions) error {
+		o.BulkInsertLimit = limit
+		return nil
+	}
+}
+
 // SQLMaxOpenConns sets the MaxOpenConns option
 func SQLMaxOpenConns(max int) SQLStoreOption {
 	return func(o *SQLStoreOptions) error {
@@ -254,6 +271,7 @@ func SQLAllOptions(opts *SQLStoreOptions) SQLStoreOption {
 	return func(o *SQLStoreOptions) error {
 		o.NoCaching = opts.NoCaching
 		o.MaxOpenConns = opts.MaxOpenConns
+		o.BulkInsertLimit = opts.BulkInsertLimit
 		return nil
 	}
 }
@@ -275,6 +293,8 @@ type SQLStore struct {
 	wg            sync.WaitGroup
 	preparedStmts []*sql.Stmt
 	ssFlusher     *subStoresFlusher
+	postgres      bool
+	bulkInserts   []string
 }
 
 type sqlDBLock struct {
@@ -410,11 +430,20 @@ func NewSQLStore(log logger.Logger, driver, source string, limits *StoreLimits,
 		db:            db,
 		doneCh:        make(chan struct{}),
 		preparedStmts: make([]*sql.Stmt, 0, len(sqlStmts)),
+		postgres:      driver == driverPostgres,
 	}
 	if err := s.init(TypeSQL, log, limits); err != nil {
 		s.Close()
 		return nil, err
 	}
+	if s.postgres && opts.BulkInsertLimit > 0 {
+		limit := opts.BulkInsertLimit
+		s.bulkInserts = make([]string, limit)
+		for i := 0; i < limit; i++ {
+			j := i * 5
+			s.bulkInserts[i] = fmt.Sprintf("($%d,$%d,$%d,$%d,$%d)", j+1, j+2, j+3, j+4, j+5)
+		}
+	}
 	if err := s.createPreparedStmts(); err != nil {
 		s.Close()
 		return nil, err
@@ -1633,6 +1662,9 @@ func (ms *SQLMsgStore) flush() error {
 		ps *sql.Stmt
 	)
 	defer func() {
+		if ms.limits.MaxAge > 0 && ms.expireTimer == nil {
+			ms.createExpireTimer()
+		}
 		ms.writeCache.transferToFreeList()
 		if ps != nil {
 			ps.Close()
@@ -1641,6 +1673,9 @@ func (ms *SQLMsgStore) flush() error {
 			tx.Rollback()
 		}
 	}()
+	if limit := ms.sqlStore.opts.BulkInsertLimit; limit >= sqlMinBulkInsertLimit {
+		return ms.bulkInsert(limit)
+	}
 	tx, err := ms.sqlStore.db.Begin()
 	if err != nil {
 		return err
@@ -1664,8 +1699,80 @@ func (ms *SQLMsgStore) flush() error {
 		return err
 	}
 	tx = nil
-	if ms.limits.MaxAge > 0 && ms.expireTimer == nil {
-		ms.createExpireTimer()
+	return nil
+}
+
+// Insert messages with INSERT INTO MESSAGES () VALUES (),(),()...
+// This is faster than the original insert with transactions.
+// It is done only if user configures the BulkInsertLimit option.
+// Lock held on entry.
+func (ms *SQLMsgStore) bulkInsert(limit int) error {
+	const insertStmt = "INSERT INTO Messages (id, seq, timestamp, size, data) VALUES "
+	const valArgs = "(?,?,?,?,?)"
+
+	count := ms.writeCache.count
+
+	sb := strings.Builder{}
+	size := len(insertStmt) + count // number of "," + last ";"
+	if ms.sqlStore.postgres {
+		for i := 0; i < limit; i++ {
+			size += len(ms.sqlStore.bulkInserts[i])
+		}
+	} else {
+		size += count * len(valArgs)
+	}
+	sb.Grow(size)
+	sb.WriteString(insertStmt)
+
+	for i := 0; i < limit; i++ {
+		if i > 0 {
+			sb.WriteString(",")
+		}
+		if ms.sqlStore.postgres {
+			sb.WriteString(ms.sqlStore.bulkInserts[i])
+		} else {
+			sb.WriteString(valArgs)
+		}
+	}
+	sb.WriteString(";")
+	stmtb := []byte(sb.String())
+
+	args := make([]interface{}, 0, 5*count)
+	start := ms.writeCache.head
+	for count > 0 {
+		args = args[:0]
+		i := 0
+		l := len(insertStmt)
+		// Iterate through the cache, but do not remove elements from the list.
+		// They are needed by the caller.
+		for cm := start; cm != nil; cm = cm.next {
+			if i > 0 {
+				l++
+			}
+			if ms.sqlStore.postgres {
+				l += len(ms.sqlStore.bulkInserts[i])
+			} else {
+				l += len(valArgs)
+			}
+			args = append(args, ms.channelID, cm.msg.Sequence, cm.msg.Timestamp, len(cm.data), cm.data)
+			i++
+			if i == limit {
+				start = cm.next
+				break
+			}
+		}
+		count -= i
+		var stmt string
+		if i == limit {
+			stmt = sb.String()
+		} else {
+			l++
+			stmtb[l-1] = ';'
+			stmt = string(stmtb[:l])
+		}
+		if _, err := ms.sqlStore.db.Exec(stmt, args[:i*5]...); err != nil {
+			return err
+		}
 	}
 	return nil
 }

stores/sqlstore_test.go

@@ -160,8 +160,9 @@ func TestSQLAllOptions(t *testing.T) {
 	defer cleanupSQLDatastore(t)
 
 	opts := &SQLStoreOptions{
-		NoCaching:    true,
-		MaxOpenConns: 123,
+		NoCaching:       true,
+		MaxOpenConns:    123,
+		BulkInsertLimit: 456,
 	}
 	s, err := NewSQLStore(testLogger, testSQLDriver, testSQLSource, nil, SQLAllOptions(opts))
 	if err != nil {
@@ -178,6 +179,9 @@ func TestSQLAllOptions(t *testing.T) {
 	if so.MaxOpenConns != 123 {
 		t.Fatalf("MaxOpenConns should be 123, got %v", so.MaxOpenConns)
 	}
+	if so.BulkInsertLimit != 456 {
+		t.Fatalf("BulkInsertLimit should be 456, got %v", so.BulkInsertLimit)
+	}
 }
 
 func TestSQLPostgresDriverInit(t *testing.T) {
@@ -2140,3 +2144,44 @@ func TestSQLMaxAgeForMsgsWithTimestampInPast(t *testing.T) {
 		}
 	}
 }
+
+func TestSQLBulkInsertLimit(t *testing.T) {
+	if !doSQL {
+		t.SkipNow()
+	}
+
+	cleanupSQLDatastore(t)
+	defer cleanupSQLDatastore(t)
+
+	// Create store with caching enabled and bulk insert limit
+	s, err := NewSQLStore(testLogger, testSQLDriver, testSQLSource, nil,
+		SQLNoCaching(false), SQLBulkInsertLimit(10))
+	if err != nil {
+		t.Fatalf("Error creating store: %v", err)
+	}
+	defer s.Close()
+
+	cs := storeCreateChannel(t, s, "foo")
+	for seq := uint64(1); seq < 127; seq++ {
+		msg := &pb.MsgProto{
+			Sequence:  seq,
+			Subject:   "foo",
+			Data:      []byte(fmt.Sprintf("%v", seq)),
+			Timestamp: time.Now().UnixNano(),
+		}
+		if _, err := cs.Msgs.Store(msg); err != nil {
+			t.Fatalf("Error storing message: %v", err)
+		}
+	}
+	if err := cs.Msgs.Flush(); err != nil {
+		t.Fatalf("Error on flush: %v", err)
+	}
+
+	for seq := uint64(1); seq < 127; seq++ {
+		m := msgStoreLookup(t, cs.Msgs, seq)
+		expected := fmt.Sprintf("%v", seq)
+		if string(m.Data) != expected {
+			t.Fatalf("Expected %q, got %q", expected, m.Data)
+		}
+	}
+}

test/configs/test_parse.conf

@@ -99,5 +99,6 @@ streaming: {
     source: "ivan:pwd@/nss_db"
     no_caching: true
     max_open_conns: 5
+    bulk_insert_limit: 1000
   }
 }