main.go

package main

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
	_ "time/tzdata"

	"github.com/google/uuid"
	_ "github.com/jackc/pgx/v5/stdlib"

	_ "embed"
)

//go:embed schema.sql
var queryCreate string

func main() {
	// runtime.GOMAXPROCS(runtime.NumCPU())

	loc, err := time.LoadLocation("UTC")
	if err != nil {
		panic(err)
	}
	time.Local = loc

	dsn := os.ExpandEnv("user=$PSQL_USERNAME password=$PSQL_PASSWORD host=localhost port=5432 dbname=$PSQL_DB sslmode=disable")
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		slog.Error("db open failed", "error", err)
		return
	}
	defer db.Close()

	_, err = db.ExecContext(context.Background(), queryCreate)
	if err != nil {
		slog.Error("db create failed", "error", err)
		return
	}
	slog.Info(`db create ok`)

	query := `INSERT INTO impressions (impression_id, ad_id, image_url, click_url) VALUES ($1, $2, $3, $4)`

	maxDuration := 1 * time.Second
	startTime := time.Now()

	var count uint64

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	timer := time.NewTimer(maxDuration)
	go func() {
		<-timer.C
		cancel()
	}()

	workerCount := runtime.NumCPU()
	var wg sync.WaitGroup
	for range workerCount {
		wg.Go(func() {
			for i := 0; ; i++ {
				select {
				case <-ctx.Done():
					return
				default:
				}

				impressionID := uuid.New().String()
				adID := "ad456"
				imageURL := "https://example.com/image.jpg"
				clickURL := "https://advertiser.com/click"

				_, err := db.ExecContext(ctx, query, impressionID, adID, imageURL, clickURL)
				if err != nil {
					slog.Error("db insert failed", `i`, i, "error", err)
					cancel()
					return
				}
				atomic.AddUint64(&count, 1)
			}
		})
	}

	wg.Wait()

	elapsed := time.Since(startTime)
	fmt.Printf("Inserted %d rows in %.2f seconds (%.2f inserts/second)\n", count, elapsed.Seconds(), float64(count)/elapsed.Seconds())
}