refactor: use batcher utility and StringVarEnv

internal/db/config.go

@@ -111,10 +111,6 @@ func (o *Options) ToConfig() (*Config, error) {
 		OmitRaw:          o.OmitRaw,
 	}
 
-	if cfg.ConnectionString == "" {
-		cfg.ConnectionString = os.Getenv(EnvConnectionString)
-	}
-
 	cfg.ApplyDefaults()
 
 	if err := cfg.Validate(); err != nil {

internal/db/writer.go

@@ -2,25 +2,21 @@ package db
 
 import (
 	"context"
-	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/httpx/runner"
+	"github.com/projectdiscovery/utils/batcher"
 )
 
 type Writer struct {
-	db       Database
-	cfg      *Config
-	data     chan runner.Result
-	done     chan struct{}
-	counter  atomic.Int64
-	closed   atomic.Bool
-	wg       sync.WaitGroup
-	ctx      context.Context
-	cancel   context.CancelFunc
-	omitRaw  bool
+	db      Database
+	cfg     *Config
+	batcher *batcher.Batcher[runner.Result]
+	counter atomic.Int64
+	closed  atomic.Bool
+	omitRaw bool
 }
 
 func NewWriter(ctx context.Context, cfg *Config) (*Writer, error) {
@@ -38,26 +34,41 @@ func NewWriter(ctx context.Context, cfg *Config) (*Writer, error) {
 		return nil, err
 	}
 
-	writerCtx, cancel := context.WithCancel(ctx)
-
 	w := &Writer{
 		db:      db,
 		cfg:     cfg,
-		data:    make(chan runner.Result, cfg.BatchSize),
-		done:    make(chan struct{}),
-		ctx:     writerCtx,
-		cancel:  cancel,
 		omitRaw: cfg.OmitRaw,
 	}
 
-	w.wg.Add(1)
-	go w.run()
+	w.batcher = batcher.New(
+		batcher.WithMaxCapacity[runner.Result](cfg.BatchSize),
+		batcher.WithFlushInterval[runner.Result](cfg.FlushInterval),
+		batcher.WithFlushCallback(w.flush),
+	)
+
+	w.batcher.Run()
 
 	gologger.Info().Msgf("Database output enabled: %s (%s/%s)", cfg.Type, cfg.DatabaseName, cfg.TableName)
 
 	return w, nil
 }
 
+func (w *Writer) flush(batch []runner.Result) {
+	if len(batch) == 0 {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	if err := w.db.InsertBatch(ctx, batch); err != nil {
+		gologger.Error().Msgf("Failed to insert batch to database: %v", err)
+	} else {
+		w.counter.Add(int64(len(batch)))
+		gologger.Verbose().Msgf("Inserted %d results to database (total: %d)", len(batch), w.counter.Load())
+	}
+}
+
 func (w *Writer) GetWriterCallback() runner.OnResultCallback {
 	return func(r runner.Result) {
 		if w.closed.Load() {
@@ -75,60 +86,7 @@ func (w *Writer) GetWriterCallback() runner.OnResultCallback {
 			r.RawHeaders = ""
 		}
 
-		select {
-		case w.data <- r:
-		case <-w.ctx.Done():
-		}
-	}
-}
-
-func (w *Writer) run() {
-	defer w.wg.Done()
-	defer close(w.done)
-
-	batch := make([]runner.Result, 0, w.cfg.BatchSize)
-	ticker := time.NewTicker(w.cfg.FlushInterval)
-	defer ticker.Stop()
-
-	flush := func() {
-		if len(batch) == 0 {
-			return
-		}
-
-		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-		defer cancel()
-
-		if err := w.db.InsertBatch(ctx, batch); err != nil {
-			gologger.Error().Msgf("Failed to insert batch to database: %v", err)
-		} else {
-			w.counter.Add(int64(len(batch)))
-			gologger.Verbose().Msgf("Inserted %d results to database (total: %d)", len(batch), w.counter.Load())
-		}
-
-		batch = batch[:0]
-	}
-
-	for {
-		select {
-		case <-w.ctx.Done():
-			flush()
-			return
-
-		case <-ticker.C:
-			flush()
-
-		case result, ok := <-w.data:
-			if !ok {
-				flush()
-				return
-			}
-
-			batch = append(batch, result)
-
-			if len(batch) >= w.cfg.BatchSize {
-				flush()
-			}
-		}
+		w.batcher.Append(r)
 	}
 }
 
@@ -137,9 +95,8 @@ func (w *Writer) Close() {
 		return
 	}
 
-	w.cancel()
-	close(w.data)
-	<-w.done
+	w.batcher.Stop()
+	w.batcher.WaitDone()
 
 	if err := w.db.Close(); err != nil {
 		gologger.Error().Msgf("Error closing database connection: %v", err)

runner/options.go

@@ -506,7 +506,7 @@ func ParseOptions() *Options {
 		flagSet.BoolVarP(&options.ResultDatabase, "result-db", "rdb", false, "store results in database"),
 		flagSet.StringVarP(&options.ResultDatabaseConfig, "result-db-config", "rdbc", "", "path to database config file"),
 		flagSet.StringVarP(&options.ResultDatabaseType, "result-db-type", "rdbt", "", "database type (mongodb, postgres, mysql)"),
-		flagSet.StringVarP(&options.ResultDatabaseConnStr, "result-db-conn", "rdbcs", "", "database connection string (env: HTTPX_DB_CONNECTION_STRING)"),
+		flagSet.StringVarEnv(&options.ResultDatabaseConnStr, "result-db-conn", "rdbcs", "", "HTTPX_DB_CONNECTION_STRING", "database connection string"),
 		flagSet.StringVarP(&options.ResultDatabaseName, "result-db-name", "rdbn", "httpx", "database name"),
 		flagSet.StringVarP(&options.ResultDatabaseTable, "result-db-table", "rdbtb", "results", "table/collection name"),
 		flagSet.IntVarP(&options.ResultDatabaseBatchSize, "result-db-batch-size", "rdbbs", 100, "batch size for database inserts"),