github.com/tachunwu/scale@v0.0.0-20230215135019-20d7de63a560/cmd/leaf/main.go (about)

     1  package main
     2  
     3  // 45.059151905 sec process 40,0000 (TPS: 8877.21990071)
     4  // 74.768570822 sec process 80,0000 (TPS: 10699.6829176)
     5  
     6  import (
     7  	"log"
     8  
     9  	"github.com/cockroachdb/pebble"
    10  	natsserver "github.com/nats-io/nats-server/v2/server"
    11  	"github.com/nats-io/nats.go"
    12  	"github.com/tachunwu/scale/pkg/database"
    13  	"github.com/tachunwu/scale/pkg/jetstream"
    14  	"go.uber.org/zap"
    15  )
    16  
    17  func main() {
    18  	app := NewLeafApp()
    19  	Client()
    20  	app.Run()
    21  }
    22  
    23  type LeafApp struct {
    24  	db     database.DB
    25  	js     *natsserver.Server
    26  	logger *zap.Logger
    27  }
    28  
    29  func NewLeafApp() *LeafApp {
    30  	logger, _ := zap.NewProduction()
    31  	defer logger.Sync()
    32  	return &LeafApp{
    33  		db:     database.NewPebbleDB("./data"),
    34  		js:     jetstream.NewJetStream(),
    35  		logger: logger,
    36  	}
    37  }
    38  
    39  func (a *LeafApp) Run() {
    40  
    41  	nc, err := nats.Connect(nats.DefaultURL)
    42  	if err != nil {
    43  		log.Fatal(err)
    44  	}
    45  	defer nc.Close()
    46  	js, err := nc.JetStream(nats.PublishAsyncMaxPending(256))
    47  	if err != nil {
    48  		a.logger.Error("JetStream error", zap.Any("error", err))
    49  		return
    50  	}
    51  
    52  	kv, err := js.KeyValue(jetstream.BucketName)
    53  	if err != nil {
    54  		a.logger.Error("KV error", zap.Any("error", err))
    55  		return
    56  	}
    57  
    58  	w, _ := kv.WatchAll()
    59  	kveCh := make(chan nats.KeyValueEntry, 1)
    60  	// Producer
    61  	go func() {
    62  		for {
    63  			kve := <-w.Updates()
    64  			kveCh <- kve
    65  		}
    66  
    67  	}()
    68  	// Consumer
    69  	for {
    70  		kve := <-kveCh
    71  		if kve != nil {
    72  			b := a.db.NewBatch()
    73  			b.Set([]byte(kve.Key()), kve.Value(), pebble.Sync)
    74  			b.Commit(pebble.Sync)
    75  			// a.logger.Info(
    76  			// 	"Pebble store",
    77  			// 	zap.String("Key", kve.Key()),
    78  			// 	zap.Binary("Value", kve.Value()),
    79  			// )
    80  		}
    81  
    82  	}
    83  }