tcp_client2.go

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"sync"
	"time"
)

func main() {
	if len(os.Args) != 5 {
		fmt.Println("Usage: go run client.go <host:port> <duration_milliseconds> <mb_per_second> <concurrent_connections>")
		fmt.Println("Example: go run client.go 127.0.0.1:9090 10000 1 5")
		os.Exit(1)
	}

	endpoint := os.Args[1]
	duration, err := strconv.Atoi(os.Args[2])
	if err != nil {
		log.Fatal("Invalid duration:", err)
	}
	mbPerSec, err := strconv.Atoi(os.Args[3])
	if err != nil {
		log.Fatal("Invalid MB per second:", err)
	}
	concurrentConns, err := strconv.Atoi(os.Args[4])
	if err != nil {
		log.Fatal("Invalid concurrent connections:", err)
	}

	if concurrentConns <= 0 {
		log.Fatal("Concurrent connections must be greater than 0")
	}

	// Calculate per-connection rate
	mbPerSecPerConn := float64(mbPerSec) / float64(concurrentConns)
	bytesPerSecPerConn := int(mbPerSecPerConn * 1024 * 1024)
	chunksPerSecPerConn := bytesPerSecPerConn / 65536
	
	if chunksPerSecPerConn == 0 {
		log.Fatal("Rate too low for the number of connections. Try reducing concurrent connections or increasing MB/sec.")
	}

	intervalBetweenChunks := time.Second / time.Duration(chunksPerSecPerConn)

	fmt.Printf("Sending %d MB/sec total (%.2f MB/sec per connection) for %d ms to %s using %d concurrent connections\n", 
		mbPerSec, mbPerSecPerConn, duration, endpoint, concurrentConns)
	fmt.Printf("Each connection: %d chunks/sec, interval: %v\n", chunksPerSecPerConn, intervalBetweenChunks)

	var wg sync.WaitGroup
	totalBytesChannel := make(chan int, concurrentConns)

	// Start concurrent connections
	for i := 0; i < concurrentConns; i++ {
		wg.Add(1)
		go func(connID int) {
			defer wg.Done()
			bytes := runConnection(endpoint, duration, intervalBetweenChunks, connID)
			totalBytesChannel <- bytes
		}(i + 1)
	}

	// Wait for all connections to finish
	wg.Wait()
	close(totalBytesChannel)

	// Calculate total bytes sent across all connections
	totalBytes := 0
	for bytes := range totalBytesChannel {
		totalBytes += bytes
	}

	fmt.Printf("Sent %d bytes total across %d connections\n", totalBytes, concurrentConns)
}

func runConnection(endpoint string, duration int, interval time.Duration, connID int) int {
	conn, err := net.Dial("tcp", endpoint)
	if err != nil {
		log.Printf("Connection %d failed: %v", connID, err)
		return 0
	}
	defer conn.Close()

	// Create 65536-byte chunk filled with 'X'
	chunk := make([]byte, 65536)
	for i := range chunk {
		chunk[i] = 'X'
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	endTime := time.Now().Add(time.Duration(duration) * time.Millisecond)
	totalBytes := 0

	for time.Now().Before(endTime) {
		<-ticker.C
		n, err := conn.Write(chunk)
		if err != nil {
			log.Printf("Connection %d write error: %v", connID, err)
			break
		}
		totalBytes += n
	}

	fmt.Printf("Connection %d sent %d bytes\n", connID, totalBytes)
	return totalBytes
}