Skip to content

Commit ade5446

Browse files
pompon0pompon0
committed
main.go wip
1 parent 1672794 commit ade5446

1 file changed

Lines changed: 157 additions & 150 deletions

File tree

‎main.go‎

Lines changed: 157 additions & 150 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package main
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"log"
78
"os"
89
"os/signal"
910
"strings"
1011
"syscall"
12+
"github.com/sei-protocol/sei-load/utils/service"
1113
"time"
1214

1315
"github.com/spf13/cobra"
@@ -73,183 +75,188 @@ func main() {
7375
}
7476

7577
func runLoadTest(cmd *cobra.Command, args []string) {
76-
// Parse the config file into a config.LoadConfig struct
77-
cfg, err := loadConfig(configFile)
78-
if err != nil {
79-
log.Fatalf("Failed to load config: %v", err)
80-
}
81-
82-
fmt.Printf("🚀 Starting Sei Chain Load Test v2\n")
83-
fmt.Printf("📁 Config file: %s\n", configFile)
84-
fmt.Printf("🎯 Endpoints: %d\n", len(cfg.Endpoints))
85-
fmt.Printf("👥 Workers per endpoint: %d\n", workers)
86-
fmt.Printf("🔧 Total workers: %d\n", len(cfg.Endpoints)*workers)
87-
fmt.Printf("📊 Scenarios: %d\n", len(cfg.Scenarios))
88-
fmt.Printf("⏱️ Stats interval: %v\n", statsInterval)
89-
fmt.Printf("📦 Buffer size per worker: %d\n", bufferSize)
90-
if tps > 0 {
91-
fmt.Printf("📈 Transactions per second: %.2f\n", tps)
92-
}
93-
if dryRun {
94-
fmt.Printf("📝 Dry run: enabled\n")
95-
}
96-
if trackReceipts {
97-
fmt.Printf("📝 Track receipts: enabled\n")
98-
}
99-
if trackBlocks {
100-
fmt.Printf("📝 Track blocks: enabled\n")
101-
}
102-
if prewarm {
103-
fmt.Printf("📝 Prewarm: enabled\n")
104-
}
105-
fmt.Println()
78+
ctx := context.Background()
79+
err := service.Run(ctx, func(ctx context.Context, s service.Scope) error {
80+
// Parse the config file into a config.LoadConfig struct
81+
cfg, err := loadConfig(configFile)
82+
if err != nil {
83+
return fmt.Errorf("Failed to load config: %w", err)
84+
}
10685

107-
// Enable mock deployment in dry-run mode
108-
if dryRun {
109-
cfg.MockDeploy = true
110-
}
86+
log.Printf("🚀 Starting Sei Chain Load Test v2")
87+
log.Printf("📁 Config file: %s", configFile)
88+
log.Printf("🎯 Endpoints: %d", len(cfg.Endpoints))
89+
log.Printf("👥 Workers per endpoint: %d", workers)
90+
log.Printf("🔧 Total workers: %d", len(cfg.Endpoints)*workers)
91+
log.Printf("📊 Scenarios: %d", len(cfg.Scenarios))
92+
log.Printf("⏱️ Stats interval: %v", statsInterval)
93+
log.Printf("📦 Buffer size per worker: %d", bufferSize)
94+
if tps > 0 {
95+
log.Printf("📈 Transactions per second: %.2f", tps)
96+
}
97+
if dryRun {
98+
log.Printf("📝 Dry run: enabled")
99+
}
100+
if trackReceipts {
101+
log.Printf("📝 Track receipts: enabled")
102+
}
103+
if trackBlocks {
104+
log.Printf("📝 Track blocks: enabled")
105+
}
106+
if prewarm {
107+
log.Printf("📝 Prewarm: enabled")
108+
}
109+
log.Println()
111110

112-
// Create the generator from the config struct
113-
gen, err := generator.NewConfigBasedGenerator(cfg)
114-
if err != nil {
115-
log.Fatalf("Failed to create generator: %v", err)
116-
}
111+
// Enable mock deployment in dry-run mode
112+
if dryRun {
113+
cfg.MockDeploy = true
114+
}
117115

118-
// Create the sender from the config struct
119-
snd, err := sender.NewShardedSender(cfg, bufferSize, workers)
120-
if err != nil {
121-
log.Fatalf("Failed to create sender: %v", err)
122-
}
116+
// Create the generator from the config struct
117+
gen, err := generator.NewConfigBasedGenerator(cfg)
118+
if err != nil {
119+
return fmt.Errorf("Failed to create generator: %w", err)
120+
}
123121

124-
// Create statistics collector and logger
125-
collector := stats.NewCollector()
126-
logger := stats.NewLogger(collector, statsInterval, debug)
127-
128-
// Create and start block collector if endpoints are available
129-
var blockCollector *stats.BlockCollector
130-
if len(cfg.Endpoints) > 0 && trackBlocks {
131-
blockCollector = stats.NewBlockCollector(cfg.Endpoints[0])
132-
collector.SetBlockCollector(blockCollector)
133-
// Start block collector
134-
if err := blockCollector.Start(); err != nil {
135-
log.Printf("⚠️ Failed to start block collector: %v", err)
122+
// Create the sender from the config struct
123+
snd, err := sender.NewShardedSender(cfg, bufferSize, workers)
124+
if err != nil {
125+
return fmt.Errorf("Failed to create sender: %w", err)
136126
}
137-
}
138127

139-
// Enable dry-run mode in sender if specified
140-
if dryRun {
141-
snd.SetDryRun(true)
142-
}
143-
if debug {
144-
snd.SetDebug(true)
145-
}
146-
if trackReceipts {
147-
snd.SetTrackReceipts(true)
148-
}
149-
if trackBlocks {
150-
snd.SetTrackBlocks(true)
151-
}
128+
// Create statistics collector and logger
129+
collector := stats.NewCollector()
130+
logger := stats.NewLogger(collector, statsInterval, debug)
131+
132+
// Create and start block collector if endpoints are available
133+
var blockCollector *stats.BlockCollector
134+
if len(cfg.Endpoints) > 0 && trackBlocks {
135+
blockCollector = stats.NewBlockCollector(cfg.Endpoints[0])
136+
collector.SetBlockCollector(blockCollector)
137+
// Start block collector
138+
if err := blockCollector.Start(); err != nil {
139+
log.Printf("⚠️ Failed to start block collector: %v", err)
140+
}
141+
}
152142

153-
// Set statistics collector for sender and its workers
154-
snd.SetStatsCollector(collector, logger)
143+
// Enable dry-run mode in sender if specified
144+
if dryRun {
145+
snd.SetDryRun(true)
146+
}
147+
if debug {
148+
snd.SetDebug(true)
149+
}
150+
if trackReceipts {
151+
snd.SetTrackReceipts(true)
152+
}
153+
if trackBlocks {
154+
snd.SetTrackBlocks(true)
155+
}
155156

156-
// Create dispatcher
157-
dispatcher := sender.NewDispatcher(gen, snd)
158-
if tps > 0 {
159-
// Convert TPS to interval: 1/tps seconds = (1/tps) * 1e9 nanoseconds
160-
intervalNs := int64((1.0 / tps) * 1e9)
161-
dispatcher.SetRateLimit(time.Duration(intervalNs))
162-
}
157+
// Set statistics collector for sender and its workers
158+
snd.SetStatsCollector(collector, logger)
163159

164-
// Set statistics collector for dispatcher
165-
dispatcher.SetStatsCollector(collector, logger)
160+
// Create dispatcher
161+
dispatcher := sender.NewDispatcher(gen, snd)
162+
if tps > 0 {
163+
// Convert TPS to interval: 1/tps seconds = (1/tps) * 1e9 nanoseconds
164+
intervalNs := int64((1.0 / tps) * 1e9)
165+
dispatcher.SetRateLimit(time.Duration(intervalNs))
166+
}
166167

167-
// Set up prewarming if enabled
168-
if prewarm {
169-
fmt.Println("🔥 Creating prewarm generator...")
170-
prewarmGen := generator.NewPrewarmGenerator(cfg, gen)
171-
dispatcher.SetPrewarmGenerator(prewarmGen)
172-
fmt.Println("✅ Prewarm generator ready")
173-
fmt.Printf("📝 Prewarm mode: Accounts will be prewarmed\n")
174-
}
168+
// Set statistics collector for dispatcher
169+
dispatcher.SetStatsCollector(collector, logger)
175170

176-
// Start the sender (starts all workers)
177-
snd.Start()
178-
fmt.Printf("✅ Connected to %d endpoints\n", snd.GetNumShards())
171+
// Set up prewarming if enabled
172+
if prewarm {
173+
fmt.Println("🔥 Creating prewarm generator...")
174+
prewarmGen := generator.NewPrewarmGenerator(cfg, gen)
175+
dispatcher.SetPrewarmGenerator(prewarmGen)
176+
log.Println("✅ Prewarm generator ready")
177+
log.Printf("📝 Prewarm mode: Accounts will be prewarmed")
178+
}
179179

180-
// Start block collector if enabled
181-
if trackBlocks {
182-
blockCollector = stats.NewBlockCollector(cfg.Endpoints[0])
183-
collector.SetBlockCollector(blockCollector)
184-
err = blockCollector.Start()
185-
if err != nil {
186-
log.Fatalf("Failed to start block collector: %v", err)
180+
// Start the sender (starts all workers)
181+
snd.Start()
182+
log.Printf("✅ Connected to %d endpoints", snd.GetNumShards())
183+
184+
// Start block collector if enabled
185+
if trackBlocks {
186+
blockCollector = stats.NewBlockCollector(cfg.Endpoints[0])
187+
collector.SetBlockCollector(blockCollector)
188+
if err := blockCollector.Start(); err != nil {
189+
return fmt.Errorf("Failed to start block collector: %w", err)
190+
}
191+
log.Println("✅ Started block collector")
187192
}
188-
fmt.Println("✅ Started block collector")
189-
}
190193

191-
// Perform prewarming if enabled (before starting logger to avoid logging prewarm transactions)
192-
if prewarm {
193-
err = dispatcher.Prewarm()
194-
if err != nil {
195-
log.Fatalf("Failed to prewarm accounts: %v", err)
194+
// Perform prewarming if enabled (before starting logger to avoid logging prewarm transactions)
195+
if prewarm {
196+
if err := dispatcher.Prewarm(); err != nil {
197+
return fmt.Errorf("Failed to prewarm accounts: %w", err)
198+
}
196199
}
197-
}
198200

199-
// Start logger (after prewarming to capture only main load test metrics)
200-
logger.Start()
201-
fmt.Println("✅ Started statistics logger")
201+
// Start logger (after prewarming to capture only main load test metrics)
202+
logger.Start()
203+
log.Println("✅ Started statistics logger")
202204

203-
// Start dispatcher for main load test
204-
dispatcher.Start()
205-
fmt.Println("✅ Started dispatcher")
205+
// Start dispatcher for main load test
206+
dispatcher.Start()
207+
log.Println("✅ Started dispatcher")
206208

207-
// Set up signal handling for graceful shutdown
208-
sigChan := make(chan os.Signal, 1)
209-
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
209+
// Set up signal handling for graceful shutdown
210+
sigChan := make(chan os.Signal, 1)
211+
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
210212

211-
fmt.Printf("📈 Logging statistics every %v (Press Ctrl+C to stop)\n", statsInterval)
212-
if dryRun {
213-
fmt.Printf("📝 Dry-run mode: Simulating requests without sending\n")
214-
}
215-
if debug {
216-
fmt.Printf("🐛 Debug mode: Each transaction will be logged\n")
217-
}
218-
if trackReceipts {
219-
fmt.Printf("📝 Track receipts mode: Receipts will be tracked\n")
220-
}
221-
if trackBlocks {
222-
fmt.Printf("📝 Track blocks mode: Block data will be collected\n")
223-
}
224-
fmt.Println(strings.Repeat("=", 60))
213+
log.Printf("📈 Logging statistics every %v (Press Ctrl+C to stop)", statsInterval)
214+
if dryRun {
215+
log.Printf("📝 Dry-run mode: Simulating requests without sending")
216+
}
217+
if debug {
218+
log.Printf("🐛 Debug mode: Each transaction will be logged")
219+
}
220+
if trackReceipts {
221+
log.Printf("📝 Track receipts mode: Receipts will be tracked")
222+
}
223+
if trackBlocks {
224+
log.Printf("📝 Track blocks mode: Block data will be collected")
225+
}
226+
fmt.Println(strings.Repeat("=", 60))
225227

226-
// Main loop - wait for shutdown signal
227-
<-sigChan
228+
// Main loop - wait for shutdown signal
229+
<-sigChan
228230

229-
fmt.Println("\n🛑 Received shutdown signal, stopping gracefully...")
231+
log.Println("🛑 Received shutdown signal, stopping gracefully...")
230232

231-
// Stop block collector first
232-
if blockCollector != nil {
233-
blockCollector.Stop()
234-
fmt.Println("✅ Stopped block collector")
235-
}
233+
// Stop block collector first
234+
if blockCollector != nil {
235+
blockCollector.Stop()
236+
log.Println("✅ Stopped block collector")
237+
}
236238

237-
// Stop statistics logger first
238-
logger.Stop()
239-
fmt.Println("✅ Stopped statistics logger")
239+
// Stop statistics logger first
240+
logger.Stop()
241+
log.Println("✅ Stopped statistics logger")
240242

241-
// Stop dispatcher
242-
dispatcher.Stop()
243-
fmt.Println("✅ Stopped dispatcher")
243+
// Stop dispatcher
244+
dispatcher.Stop()
245+
log.Println("✅ Stopped dispatcher")
244246

245-
// Stop sender and all workers
246-
snd.Stop()
247-
fmt.Println("✅ Stopped sender and workers")
247+
// Stop sender and all workers
248+
snd.Stop()
249+
log.Println("✅ Stopped sender and workers")
248250

249-
// Print final statistics
250-
logger.LogFinalStats()
251+
// Print final statistics
252+
logger.LogFinalStats()
251253

252-
fmt.Println("👋 Shutdown complete")
254+
log.Println("👋 Shutdown complete")
255+
return nil
256+
})
257+
if err != nil {
258+
log.Fatal(err)
259+
}
253260
}
254261

255262
// loadConfig reads and parses the configuration file

0 commit comments

Comments
 (0)