rtdb_bench_2.go

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"runtime"
	"time"

	"firebase.google.com/go"
)

type Write struct {
	key  string
	data map[string]interface{}
}

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 = 500
	defaultTransportPointer.MaxIdleConnsPerHost = 500

	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)
	}

	requests := make(chan *Write)
	worker := func(id int) {
		for {
			requests <- &Write{
				key: fmt.Sprintf("bench/key%d", id),
				data: map[string]interface{}{
					"online":    true,
					"timestamp": time.Now().Unix(),
				},
			}
			time.Sleep(time.Second * 10)
		}
	}

	events := make(chan bool)
	processor := func() {
		for {
			select {
			case req := <-requests:
				ref := client.NewRef(req.key)
				if err := ref.Set(ctx, req.key); err != nil {
					log.Printf("Error %s: %v", req.key, err)
				}
				events <- true
			}
		}
	}
	for i := 0; i < 10000; i++ {
		id := i
		go worker(id)
	}

	for i := 0; i < 128; i++ {
		go processor()
	}

	total := 0
	m := &runtime.MemStats{}
	for {
		select {
		case <-events:
			total++
			if total > 0 && total%500 == 0 {
				log.Print("Completed writes: ", total)
			}
			if total > 0 && total%5000 == 0 {
				runtime.ReadMemStats(m)
				log.Print("Memory usage: ", (m.Alloc / 1024 / 1024))
			}
		}
	}
}