main.go

package main

import (
	"context"
	"fmt"
	"math/rand"
	"os"
	"sync"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	databaseName     = "testdb"
	collectionName   = "testcollection"
	maxDocuments     = 100000
	requestsPerSecond = 100
	parallelClients  = 50
)

func main() {
	// Retrieve MongoDB credentials and port from the environment
	username := os.Getenv("MONGO_USERNAME")
	password := os.Getenv("MONGO_PASSWORD")
	port := os.Getenv("MONGO_PORT")
	if username == "" || password == "" {
		fmt.Println("Error: MONGO_USERNAME and MONGO_PASSWORD environment variables must be set")
		return
	}
	if port == "" {
		port = "27017" // Default MongoDB port
	}

	// Construct the MongoDB URI
	mongoURI := fmt.Sprintf("mongodb://%s:%s@localhost:%s", username, password, port)

	// Create MongoDB client
	clientOptions := options.Client().ApplyURI(mongoURI)
	client, err := mongo.Connect(context.Background(), clientOptions)
	if err != nil {
		fmt.Printf("Failed to connect to MongoDB: %v\n", err)
		return
	}
	defer client.Disconnect(context.Background())

	// Check if the database exists
	databaseNames, err := client.ListDatabaseNames(context.Background(), bson.D{})
	if err != nil {
		fmt.Printf("Failed to list databases: %v\n", err)
		return
	}

	databaseExists := false
	for _, name := range databaseNames {
		if name == databaseName {
			databaseExists = true
			break
		}
	}

	if !databaseExists {
		fmt.Printf("Database '%s' does not exist. It will be created automatically upon first write.\n", databaseName)
	}

	// Get the database
	database := client.Database(databaseName)

	// Check if the collection exists, and create it if it does not
	collectionNames, err := database.ListCollectionNames(context.Background(), bson.D{})
	if err != nil {
		fmt.Printf("Failed to list collections: %v\n", err)
		return
	}

	collectionExists := false
	for _, name := range collectionNames {
		if name == collectionName {
			collectionExists = true
			break
		}
	}

	if !collectionExists {
		err := database.CreateCollection(context.Background(), collectionName)
		if err != nil {
			fmt.Printf("Failed to create collection: %v\n", err)
			return
		}
		fmt.Printf("Collection '%s' created successfully.\n", collectionName)
	}

	// Get the collection
	collection := database.Collection(collectionName)

	// Synchronization
	var wg sync.WaitGroup
	var mu sync.Mutex
	totalRequests := 0
	startTime := time.Now()

	// Function to insert random documents
	insertRandomDocuments := func() {
		defer wg.Done()
		ticker := time.NewTicker(time.Second / requestsPerSecond)
		defer ticker.Stop()

		for {
			<-ticker.C
			mu.Lock()
			if totalRequests >= maxDocuments {
				mu.Unlock()
				break
			}
			totalRequests++
			mu.Unlock()

			// Create a random document
			doc := bson.M{
				"key":   rand.Intn(1000000),
				"value": randString(10),
			}

			// Insert the document
			_, err := collection.InsertOne(context.Background(), doc)
			if err != nil {
				fmt.Printf("Failed to insert document: %v\n", err)
			}
		}
	}

	// Start monitoring RPS
	go func() {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()

		for {
			<-ticker.C
			mu.Lock()
			if totalRequests >= maxDocuments {
				mu.Unlock()
				break
			}
			elapsed := time.Since(startTime).Seconds()
			fmt.Printf("Current RPS: %.2f\n", float64(totalRequests)/elapsed)
			mu.Unlock()
		}
	}()

	// Start parallel workers
	for i := 0; i < parallelClients; i++ {
		wg.Add(1)
		go insertRandomDocuments()
	}

	// Wait for all workers to finish
	wg.Wait()

	// Print overall RPS
	elapsed := time.Since(startTime).Seconds()
	fmt.Printf("Finished writing %d documents. Overall RPS: %.2f\n", totalRequests, float64(totalRequests)/elapsed)
}

// Generate a random string of the given length
func randString(length int) string {
	chars := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	result := make([]byte, length)
	for i := range result {
		result[i] = chars[rand.Intn(len(chars))]
	}
	return string(result)
}