fix: batch insert logs (#48)

* test: add benchmark for session log writer

* feat: implement batch import

* chore: remove debug fmt.Println

Author: bcho

server/internal/db/dbext/parse.go

@@ -16,7 +16,7 @@ type kqlParse struct {
 
 // Parse implements the ku_parse function.
 //
-//   select json_extract(ku_parse(lines, '.*'), '$.foo') as foo from source
+//	select json_extract(ku_parse(lines, '.*'), '$.foo') as foo from source
 func (p *kqlParse) Parse(fieldValue string, regexpPattern string) (string, error) {
 	re, err := regexp.Compile("(?m)" + regexpPattern)
 	if err != nil {

server/internal/db/session_io.go

@@ -9,6 +9,17 @@ import (
 	"time"
 )
 
+// TODO(hbc): tune based on usage?
+const (
+	sessionLogBuffer       = 500
+	sessionLogSendInterval = 50 * time.Millisecond
+)
+
+type sendSessionLog struct {
+	Line    string
+	ScanErr error
+}
+
 type sessionLogWriter struct {
 	Session
 
@@ -17,12 +28,15 @@ type sessionLogWriter struct {
 
 	scannerInput io.WriteCloser
 	scanner      *bufio.Scanner
-	scannerErr   error
 	scanning     *sync.WaitGroup
+
+	sendLogsChan chan sendSessionLog
+	sending      *sync.WaitGroup
+	sendErr      error
 }
 
 // SessionLogWriteCloser wraps the session as a io.WriteCloser instance.
-// The behavior of concurrency write to this writer is undefined.
+// It's not allowed to call Write() concurrently.
 func SessionLogWriteCloser(
 	rootCtx context.Context,
 	session Session,
@@ -40,59 +54,114 @@ func SessionLogWriteCloser(
 		scannerInput: pw,
 		scanner:      bufio.NewScanner(pr),
 		scanning:     new(sync.WaitGroup),
+
+		sendLogsChan: make(chan sendSessionLog, sessionLogBuffer),
+		sending:      new(sync.WaitGroup),
 	}
 
 	rv.scanning.Add(1)
 	go rv.scan()
 
+	rv.sending.Add(1)
+	go rv.sendLogs()
+
 	return rv
 }
 
 var _ io.WriteCloser = (*sessionLogWriter)(nil)
 
-func (swr *sessionLogWriter) writeLogLine(logLine string) error {
-	ctx, cancel := swr.createCtx()
-	defer cancel()
+func (swr *sessionLogWriter) sendLogs() {
+	defer swr.sending.Done()
 
-	err := swr.Session.WriteLogLine(
-		ctx,
-		WriteLogLinePayload{
+	sendTicker := time.NewTicker(sessionLogSendInterval)
+	defer sendTicker.Stop()
+
+	sendLogs := func(lines []string) error {
+		if len(lines) == 0 {
+			return nil
+		}
+
+		ctx, cancel := swr.createCtx()
+		defer cancel()
+
+		err := swr.Session.WriteLogLinesBatch(ctx, WriteLogLinesBatchPayload{
+			Lines:     lines,
 			Timestamp: swr.nowFn(),
-			Line:      logLine,
-		},
-	)
-	if err != nil {
-		return fmt.Errorf("WriteLogLine: %w", err)
+		})
+		if err != nil {
+			return fmt.Errorf("WriteLogLinesBatch: %w", err)
+		}
+
+		return nil
 	}
 
-	return nil
+	var lines []string
+	for {
+		select {
+		case sendLog, ok := <-swr.sendLogsChan:
+			if !ok || sendLog.ScanErr != nil {
+				// channel closed or scan failed, send all logs and store error
+				sendErr := sendLogs(lines)
+				if sendErr != nil {
+					swr.sendErr = sendErr
+				} else if sendLog.ScanErr != nil {
+					swr.sendErr = sendLog.ScanErr
+				}
+				return
+			}
+
+			// append logs to send buffer
+			lines = append(lines, sendLog.Line)
+		case <-sendTicker.C:
+			// send ticker fired, send all buffered logs
+			if err := sendLogs(lines); err != nil {
+				swr.sendErr = err
+				return
+			}
+			lines = nil
+		}
+	}
 }
 
 func (swr *sessionLogWriter) scan() {
 	defer swr.scanning.Done()
 
 	for swr.scanner.Scan() {
-		if err := swr.writeLogLine(swr.scanner.Text()); err != nil {
-			swr.scannerErr = err
-			return
+		swr.sendLogsChan <- sendSessionLog{
+			Line: swr.scanner.Text(),
 		}
 	}
 
 	if err := swr.scanner.Err(); err != nil {
-		swr.scannerErr = err
+		swr.sendLogsChan <- sendSessionLog{
+			ScanErr: err,
+		}
 	}
 }
 
 func (swr *sessionLogWriter) Write(p []byte) (int, error) {
+	// FIXME(hbc): race on the sendErr
+	if swr.sendErr != nil {
+		return 0, swr.sendErr
+	}
+
 	return swr.scannerInput.Write(p)
 }
 
 func (swr *sessionLogWriter) Close() error {
+	// stop scanner
 	if err := swr.scannerInput.Close(); err != nil {
 		return err
 	}
-
-	// wait for scanner to finish
 	swr.scanning.Wait()
-	return swr.scannerErr
+
+	// stop logs sender
+	close(swr.sendLogsChan)
+	swr.sending.Wait()
+
+	if swr.sendErr != nil {
+		return swr.sendErr
+	}
+
+	return nil
 }

server/internal/db/session_io_test.go

@@ -2,23 +2,24 @@ package db
 
 import (
 	"context"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
 type mockSession struct {
-	WriteLogLineFuc func(ctx context.Context, payload WriteLogLinePayload) error
+	WriteLogLinesBatchFunc func(ctx context.Context, payload WriteLogLinesBatchPayload) error
 }
 
 var _ Session = (*mockSession)(nil)
 
-func (s *mockSession) WriteLogLine(
+func (s *mockSession) WriteLogLinesBatch(
 	ctx context.Context,
-	payload WriteLogLinePayload,
+	payload WriteLogLinesBatchPayload,
 ) error {
-	return s.WriteLogLineFuc(ctx, payload)
+	return s.WriteLogLinesBatchFunc(ctx, payload)
 }
 
 func TestSessionLogWriter(t *testing.T) {
@@ -32,11 +33,11 @@ func TestSessionLogWriter(t *testing.T) {
 	}
 
 	t.Run("write serial", func(t *testing.T) {
-		var wrote []WriteLogLinePayload
+		var wrote []string
 
 		mockSession := &mockSession{
-			WriteLogLineFuc: func(ctx context.Context, payload WriteLogLinePayload) error {
-				wrote = append(wrote, payload)
+			WriteLogLinesBatchFunc: func(ctx context.Context, payload WriteLogLinesBatchPayload) error {
+				wrote = append(wrote, payload.Lines...)
 
 				return nil
 			},
@@ -53,16 +54,16 @@ func TestSessionLogWriter(t *testing.T) {
 		require.NoError(t, w.Close())
 
 		require.Len(t, wrote, 2)
-		require.Equal(t, "hello", wrote[0].Line)
-		require.Equal(t, "world", wrote[1].Line)
+		require.Equal(t, "hello", wrote[0])
+		require.Equal(t, "world", wrote[1])
 	})
 
 	t.Run("write partial", func(t *testing.T) {
-		var wrote []WriteLogLinePayload
+		var wrote []string
 
 		mockSession := &mockSession{
-			WriteLogLineFuc: func(ctx context.Context, payload WriteLogLinePayload) error {
-				wrote = append(wrote, payload)
+			WriteLogLinesBatchFunc: func(ctx context.Context, payload WriteLogLinesBatchPayload) error {
+				wrote = append(wrote, payload.Lines...)
 
 				return nil
 			},
@@ -81,8 +82,39 @@ func TestSessionLogWriter(t *testing.T) {
 		require.NoError(t, w.Close())
 
 		require.Len(t, wrote, 3)
-		require.Equal(t, "hello", wrote[0].Line)
-		require.Equal(t, "world", wrote[1].Line)
-		require.Equal(t, "foo", wrote[2].Line)
+		require.Equal(t, "hello", wrote[0])
+		require.Equal(t, "world", wrote[1])
+		require.Equal(t, "foo", wrote[2])
 	})
 }
+
+// go test -run=^$ -benchtime 30s -bench ^BenchmarkSessionLogWriter$ github.com/b4fun/ku/server/internal/db
+func BenchmarkSessionLogWriter(b *testing.B) {
+	run := func() {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		p := b.TempDir()
+		dbFile := filepath.Join(p, "test.db")
+		dbProvider, err := NewSqliteProvider(dbFile)
+		require.NoError(b, err)
+
+		_, session, err := dbProvider.CreateSession(ctx, &CreateSessionOpts{Prefix: "test"})
+		require.NoError(b, err)
+
+		writer := SessionLogWriteCloser(ctx, session, 100*time.Millisecond)
+
+		for i := 0; i < 10000; i++ {
+			content := []byte("hello world\n")
+
+			_, err := writer.Write(content)
+			require.NoError(b, err)
+		}
+
+		require.NoError(b, writer.Close())
+	}
+
+	for i := 0; i < b.N; i++ {
+		run()
+	}
+}

server/internal/db/sqlite_session.go

@@ -79,11 +79,19 @@ func (s *SqliteSession) dbTableName(tableName string) string {
 	return fmt.Sprintf("%s_%s", s.sessionID, tableName)
 }
 
-func (s *SqliteSession) WriteLogLine(ctx context.Context, payload WriteLogLinePayload) error {
-	const insertTmpl = `INSERT INTO %s (ts, lines) VALUES (?, ?)`
+func (s *SqliteSession) WriteLogLinesBatch(ctx context.Context, payload WriteLogLinesBatchPayload) error {
+	const insertTmpl = `INSERT INTO %s (ts, lines) VALUES (:ts, :lines)`
+
+	insertPayload := make([]map[string]interface{}, len(payload.Lines))
+	for i, line := range payload.Lines {
+		insertPayload[i] = map[string]interface{}{
+			"ts":    payload.Timestamp,
+			"lines": line,
+		}
+	}
 
 	stmt := fmt.Sprintf(insertTmpl, s.dbTableName(tableNameRaw))
-	if _, err := s.db.ExecContext(ctx, stmt, payload.Timestamp, payload.Line); err != nil {
+	if _, err := s.db.NamedExecContext(ctx, stmt, insertPayload); err != nil {
 		return fmt.Errorf("failed to insert log line: %w", err)
 	}
 

server/internal/db/sqlite_test.go

@@ -58,14 +58,14 @@ func TestSqliteProvider(t *testing.T) {
 		require.NotEmpty(tc, sessionID)
 		require.NotNil(tc, session)
 
-		err = session.WriteLogLine(ctx, WriteLogLinePayload{
+		err = session.WriteLogLinesBatch(ctx, WriteLogLinesBatchPayload{
 			Timestamp: time.Now(),
-			Line:      "hello",
+			Lines:     []string{"hello"},
 		})
 		require.NoError(tc, err)
-		err = session.WriteLogLine(ctx, WriteLogLinePayload{
+		err = session.WriteLogLinesBatch(ctx, WriteLogLinesBatchPayload{
 			Timestamp: time.Now(),
-			Line:      "world",
+			Lines:     []string{"world"},
 		})
 		require.NoError(tc, err)
 		sqliteSession := session.(*SqliteSession)
@@ -119,7 +119,6 @@ func TestSqliteProvider(t *testing.T) {
 		require.NotNil(tc, sessionUpdated)
 		require.Len(tc, sessionUpdated.Tables, 2)
 		for _, table := range sessionUpdated.Tables {
-			fmt.Println(table.Columns)
 			require.NotEmpty(tc, table.Columns)
 		}
 	})

server/internal/db/types.go

@@ -7,13 +7,13 @@ import (
 	v1 "github.com/b4fun/ku/protos/api/v1"
 )
 
-type WriteLogLinePayload struct {
+type WriteLogLinesBatchPayload struct {
 	Timestamp time.Time
-	Line      string
+	Lines     []string
 }
 
 type Session interface {
-	WriteLogLine(ctx context.Context, payload WriteLogLinePayload) error
+	WriteLogLinesBatch(ctx context.Context, payload WriteLogLinesBatchPayload) error
 }
 
 type CreateSessionOpts struct {