rtdb_bench.go

package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"firebase.google.com/go"
)

func main() {
	// Customize the Transport to have larger connection pool
	defaultRoundTripper := http.DefaultTransport
	defaultTransportPointer, ok := defaultRoundTripper.(*http.Transport)
	if !ok {
		panic(fmt.Sprintf("defaultRoundTripper not an *http.Transport"))
	}
	defaultTransportPointer.MaxIdleConns = 250
	defaultTransportPointer.MaxIdleConnsPerHost = 250

	ctx := context.Background()
	app, err := firebase.NewApp(ctx, &firebase.Config{
		DatabaseURL: "https://dbname.firebaseio.com",
	})
	if err != nil {
		log.Fatalln(err)
	}
	client, err := app.Database(ctx)
	if err != nil {
		log.Fatalln(err)
	}

	wg := sync.WaitGroup{}
	failures := 0
	queue := make(chan int)
	handleWrite := func(id int) {
		wg.Add(1)
		defer wg.Done()
		events := 0
		for {
			select {
			case val := <-queue:
				if val == 0 {
					log.Println("Exiting")
					return
				}
				ref := client.NewRef(fmt.Sprintf("testkey%d", val))
				if err := ref.Set(ctx, rand.Intn(1000)); err != nil {
					log.Printf("Error in goroutine %d: %v", id, err)
					failures++
				}
				events++
			}
		}
	}

	start := time.Now()
	for i := 0; i < 48; i++ {
		go handleWrite(i)
	}
	for i := 0; i < 10000; i++ {
		queue <- (rand.Intn(10) + 1)
		if i > 0 && i%500 == 0 {
			log.Printf("Sent event %d", i)
		}
	}
	close(queue)
	wg.Wait()
	log.Println("Time elapsed", time.Since(start))
	log.Println("Total failures", failures)
}