Cut per-request DB work on osctrl-tls and osctrl-api hot paths

cmd/api/auth.go

@@ -5,6 +5,7 @@ import (
 	"crypto/subtle"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/jmpsec/osctrl/cmd/api/handlers"
 	"github.com/jmpsec/osctrl/pkg/config"
@@ -156,7 +157,7 @@ func handlerAuthCheck(h http.Handler, auth, jwtSecret string) http.Handler {
 				return
 			}
 			// Update metadata for the user
-			if err := apiUsers.UpdateTokenIPAddress(utils.GetIP(r), claims.Username); err != nil {
+			if err := apiUsers.RecordTokenUse(user, utils.GetIP(r), time.Now()); err != nil {
 				log.Err(err).Msgf("error updating token for user %s", claims.Username)
 			}
 			// Set middleware values

cmd/api/handlers/stats.go

@@ -15,7 +15,6 @@ import (
 	"github.com/jmpsec/osctrl/pkg/dbutil"
 	"github.com/jmpsec/osctrl/pkg/logging"
 	"github.com/jmpsec/osctrl/pkg/nodes"
-	"github.com/jmpsec/osctrl/pkg/queries"
 	"github.com/jmpsec/osctrl/pkg/settings"
 	"github.com/jmpsec/osctrl/pkg/users"
 	"github.com/jmpsec/osctrl/pkg/utils"
@@ -107,11 +106,22 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) {
 		Environments:  make([]EnvStats, 0, len(allEnvs)),
 	}
 
+	// Active query and carve totals for every environment in one grouped
+	// count. If it fails, every environment is left out, as the per-env
+	// lookups it replaces skipped any environment whose count failed.
+	activeCounts, countErr := h.Queries.ActiveCountsByEnvironment()
+	if countErr != nil {
+		log.Warn().Err(countErr).Msg("stats: failed to count active queries and carves, skipping environments")
+	}
+
 	for _, e := range allEnvs {
 		// Filter to envs the user can actually see.
 		if !h.Users.CheckPermissions(user, users.UserLevel, e.UUID) {
 			continue
 		}
+		if countErr != nil {
+			continue
+		}
 
 		envHours := h.Settings.InactiveHours(e.ID)
 		ns, err := h.Nodes.GetStatsByEnv(e.Name, envHours)
@@ -129,19 +139,7 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) {
 			log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to get platform counts, defaulting to zeros")
 		}
 
-		// Use type-specific methods to avoid double-counting:
-		//   GetQueries returns StandardQueryType active items only.
-		//   GetCarves  returns CarveQueryType active items only.
-		activeQ, err := h.Queries.GetQueries(queries.TargetActive, e.ID)
-		if err != nil {
-			log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to count active queries, skipping env")
-			continue
-		}
-		activeC, err := h.Queries.GetCarves(queries.TargetActive, e.ID)
-		if err != nil {
-			log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to count active carves, skipping env")
-			continue
-		}
+		active := activeCounts[e.ID]
 
 		row := EnvStats{
 			InactiveHours:  envHours,
@@ -150,16 +148,16 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) {
 			Active:         ns.Active,
 			Inactive:       ns.Inactive,
 			Total:          ns.Total,
-			ActiveQueries:  len(activeQ),
-			ActiveCarves:   len(activeC),
+			ActiveQueries:  active.Queries,
+			ActiveCarves:   active.Carves,
 			PlatformCounts: platCounts,
 		}
 		out.Environments = append(out.Environments, row)
 		out.ActiveNodes += ns.Active
 		out.InactiveNodes += ns.Inactive
 		out.TotalNodes += ns.Total
-		out.TotalActiveQueries += len(activeQ)
-		out.TotalActiveCarves += len(activeC)
+		out.TotalActiveQueries += active.Queries
+		out.TotalActiveCarves += active.Carves
 		// Aggregate cross-env platform totals.
 		out.PlatformCounts.Linux += platCounts.Linux
 		out.PlatformCounts.Darwin += platCounts.Darwin
@@ -923,18 +921,17 @@ func (h *HandlersApi) NodeActivityTilesBatchHandler(w http.ResponseWriter, r *ht
 		return
 	}
 
-	// Only include nodes that actually belong to this env.
+	// Only include nodes that actually belong to this env. A lookup failure
+	// leaves nodes out rather than failing the poll, as the per-node lookups
+	// it replaces did.
 	out := make(map[string]activity.NodeTileSeries, len(uuids))
-	for _, u := range uuids {
-		node, err := h.Nodes.GetByUUID(u)
-		if err != nil {
-			continue
-		}
-		if !strings.EqualFold(node.Environment, env.Name) {
-			continue
-		}
-		if s, ok := series[node.UUID]; ok {
-			out[node.UUID] = s
+	inEnv, err := h.Nodes.UUIDsInEnvironment(uuids, env.Name)
+	if err != nil {
+		log.Warn().Err(err).Str("env", env.Name).Msg("activity tiles: node membership lookup failed")
+	}
+	for u := range inEnv {
+		if s, ok := series[u]; ok {
+			out[u] = s
 		}
 	}
 

cmd/api/handlers/stats_counts_test.go

@@ -0,0 +1,54 @@
+package handlers
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/jmpsec/osctrl/pkg/auditlog"
+	"github.com/jmpsec/osctrl/pkg/config"
+	"github.com/jmpsec/osctrl/pkg/queries"
+	"github.com/jmpsec/osctrl/pkg/settings"
+)
+
+// The stats endpoint reports active queries and carves per environment from
+// one grouped count. Console and file-explorer rows share the table and must
+// not be counted, nor finished queries.
+func TestStatsHandlerActiveQueryAndCarveCounts(t *testing.T) {
+	db, h, env, _ := setupConsoleHandlers(t)
+	h.DebugHTTPConfig = &config.YAMLConfigurationDebug{}
+	var err error
+	h.AuditLog, err = auditlog.CreateAuditLogManager(db, config.ServiceAPI, false)
+	require.NoError(t, err)
+	require.NoError(t, h.Settings.NewIntegerValue(config.ServiceAPI, settings.InactiveHours, 24, 0))
+
+	for i, dq := range []queries.DistributedQuery{
+		{Type: queries.StandardQueryType, Active: true},
+		{Type: queries.StandardQueryType, Active: true},
+		{Type: queries.StandardQueryType, Active: true, Completed: true},
+		{Type: queries.CarveQueryType, Active: true},
+		{Type: queries.ConsoleQueryType, Active: true},
+		{Type: queries.FileExplorerQueryType, Active: true},
+	} {
+		dq.Name = fmt.Sprintf("stats_q%d", i)
+		dq.Query = "SELECT 1;"
+		dq.EnvironmentID = env.ID
+		require.NoError(t, db.Create(&dq).Error)
+	}
+
+	rr := httptest.NewRecorder()
+	h.StatsHandler(rr, consoleRequest(http.MethodGet, "/stats", nil, "alice"))
+	require.Equal(t, http.StatusOK, rr.Code, rr.Body.String())
+
+	var stats StatsResponse
+	require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &stats))
+	require.Len(t, stats.Environments, 1)
+	require.Equal(t, 2, stats.Environments[0].ActiveQueries)
+	require.Equal(t, 1, stats.Environments[0].ActiveCarves)
+	require.Equal(t, 2, stats.TotalActiveQueries)
+	require.Equal(t, 1, stats.TotalActiveCarves)
+}

pkg/logging/db.go

@@ -148,21 +148,21 @@ func (logDB *LoggerDB) Status(data []byte, environment, uuid string, debug bool)
 	if err := json.Unmarshal(data, &logs); err != nil {
 		log.Err(err).Msgf("error parsing logs %s %v", string(data), err)
 	}
-	// Iterate and insert in DB
+	stamp := logDB.rowStamper()
+	entries := make([]OsqueryStatusData, 0, len(logs))
 	for _, l := range logs {
-		entry := OsqueryStatusData{
+		entries = append(entries, OsqueryStatusData{
+			Model:       stamp(),
 			UUID:        strings.ToUpper(l.HostIdentifier),
 			Environment: environment,
 			Line:        strconv.Itoa(int(l.Line)),
 			Message:     l.Message,
 			Version:     l.Version,
 			Filename:    l.Filename,
 			Severity:    strconv.Itoa(int(l.Severity)),
-		}
-		if err := logDB.Database.Conn.Create(&entry).Error; err != nil {
-			log.Err(err).Msg("Error creating status log entry")
-		}
+		})
 	}
+	insertLogRows(logDB.Database.Conn, entries, "status")
 }
 
 // Result - Function that sends JSON result logs to the configured DB
@@ -172,20 +172,82 @@ func (logDB *LoggerDB) Result(data []byte, environment, uuid string, debug bool)
 	if err != nil {
 		log.Err(err).Msgf("error parsing logs %s", string(data))
 	}
-	// Iterate and insert in DB
+	stamp := logDB.rowStamper()
+	entries := make([]OsqueryResultData, 0, len(logs))
 	for _, l := range logs {
-		entry := OsqueryResultData{
+		entries = append(entries, OsqueryResultData{
+			Model:       stamp(),
 			UUID:        strings.ToUpper(l.HostIdentifier),
 			Environment: environment,
 			Name:        l.Name,
 			Action:      l.Action,
 			Epoch:       l.Epoch,
 			Columns:     string(l.Columns),
 			Counter:     l.Counter,
+		})
+	}
+	insertLogRows(logDB.Database.Conn, entries, "result")
+}
+
+// logInsertBatchSize caps rows per multi-row INSERT. Both log models bind 10
+// columns per row (the gorm.Model timestamps plus seven fields), so 50 rows is
+// 500 bind parameters: under the 999 limit of older SQLite builds and far
+// under the 65535 of Postgres and MySQL.
+const logInsertBatchSize = 50
+
+// insertLogRows writes one osquery log batch with multi-row INSERTs inside a
+// single transaction.
+//
+// Row-by-row inserts each ran in GORM's implicit per-statement transaction,
+// so a batch of N lines cost N commits — N WAL flushes on Postgres — on the
+// busiest write path in osctrl-tls. One transaction makes that one commit.
+//
+// The batch is all-or-nothing, which row-by-row was not: one unstorable row
+// (Postgres rejects NUL bytes in text, for instance) used to cost only itself.
+// So on failure the rolled-back batch is retried row by row from a pristine
+// copy, keeping the old partial-success behaviour without duplicating
+// anything — the transaction guarantees none of the batch was kept.
+func insertLogRows[T any](db *gorm.DB, entries []T, kind string) {
+	if len(entries) == 0 {
+		return
+	}
+	// CreateInBatches writes generated IDs back into the rows it inserted,
+	// even in a batch that is later rolled back. Retrying those rows would
+	// insert explicit IDs, so the fallback works from an untouched copy.
+	pristine := append([]T(nil), entries...)
+	err := db.Transaction(func(tx *gorm.DB) error {
+		return tx.CreateInBatches(entries, logInsertBatchSize).Error
+	})
+	if err == nil {
+		return
+	}
+	log.Warn().Err(err).Int("rows", len(entries)).Msgf("batched %s log insert failed, retrying row by row", kind)
+	for i := range pristine {
+		if err := db.Create(&pristine[i]).Error; err != nil {
+			log.Err(err).Msgf("Error creating %s log entry", kind)
 		}
-		if err := logDB.Database.Conn.Create(&entry).Error; err != nil {
-			log.Err(err).Msg("Error creating result log entry")
-		}
+	}
+}
+
+// rowStamper returns a function handing out strictly increasing timestamps,
+// one per row of a log batch.
+//
+// GORM stamps every row of a multi-row INSERT with the same created_at, but
+// the log readers order by created_at alone, so ties would let the lines of
+// one batch come back shuffled. Row-by-row inserts never tied in practice;
+// stepping by the backend's storage precision (MySQL keeps milliseconds,
+// Postgres microseconds) keeps every row distinct after the database rounds
+// it, and preserves the order osquery sent the lines in.
+func (logDB *LoggerDB) rowStamper() func() gorm.Model {
+	step := time.Microsecond
+	if logDB.Database.Conn.Name() == "mysql" {
+		step = time.Millisecond
+	}
+	next := time.Now()
+	return func() gorm.Model {
+		ts := next
+		next = next.Add(step)
+		return gorm.Model{CreatedAt: ts, UpdatedAt: ts}
 	}
 }