github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/cmd/swarm-bench/benchmark.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"strconv"
     8  	"time"
     9  
    10  	"github.com/docker/swarmkit/api"
    11  	"google.golang.org/grpc"
    12  )
    13  
    14  // Config holds the benchmarking configuration.
    15  type Config struct {
    16  	Count   uint64
    17  	Manager string
    18  	IP      string
    19  	Port    int
    20  	Unit    time.Duration
    21  }
    22  
    23  // Benchmark represents a benchmark session.
    24  type Benchmark struct {
    25  	cfg       *Config
    26  	collector *Collector
    27  }
    28  
    29  // NewBenchmark creates a new benchmark session with the given configuration.
    30  func NewBenchmark(cfg *Config) *Benchmark {
    31  	return &Benchmark{
    32  		cfg:       cfg,
    33  		collector: NewCollector(),
    34  	}
    35  }
    36  
    37  // Run starts the benchmark session and waits for it to be completed.
    38  func (b *Benchmark) Run(ctx context.Context) error {
    39  	fmt.Printf("Listening for incoming connections at %s:%d\n", b.cfg.IP, b.cfg.Port)
    40  	if err := b.collector.Listen(b.cfg.Port); err != nil {
    41  		return err
    42  	}
    43  	j, err := b.launch(ctx)
    44  	if err != nil {
    45  		return err
    46  	}
    47  	fmt.Printf("Service %s launched (%d instances)\n", j.ID, b.cfg.Count)
    48  
    49  	// Periodically print stats.
    50  	doneCh := make(chan struct{})
    51  	go func() {
    52  		for {
    53  			select {
    54  			case <-time.After(5 * time.Second):
    55  				fmt.Printf("\n%s: Progression report\n", time.Now())
    56  				b.collector.Stats(os.Stdout, time.Second)
    57  			case <-doneCh:
    58  				return
    59  			}
    60  		}
    61  	}()
    62  
    63  	fmt.Println("Collecting metrics...")
    64  	b.collector.Collect(ctx, b.cfg.Count)
    65  	doneCh <- struct{}{}
    66  
    67  	fmt.Printf("\n%s: Benchmark completed\n", time.Now())
    68  	b.collector.Stats(os.Stdout, time.Second)
    69  
    70  	return nil
    71  }
    72  
    73  func (b *Benchmark) spec() *api.ServiceSpec {
    74  	return &api.ServiceSpec{
    75  		Annotations: api.Annotations{
    76  			Name: "benchmark",
    77  		},
    78  		Task: api.TaskSpec{
    79  			Runtime: &api.TaskSpec_Container{
    80  				Container: &api.ContainerSpec{
    81  					Image:   "alpine:latest",
    82  					Command: []string{"nc", b.cfg.IP, strconv.Itoa(b.cfg.Port)},
    83  				},
    84  			},
    85  		},
    86  		Mode: &api.ServiceSpec_Replicated{
    87  			Replicated: &api.ReplicatedService{
    88  				Replicas: b.cfg.Count,
    89  			},
    90  		},
    91  	}
    92  }
    93  
    94  func (b *Benchmark) launch(ctx context.Context) (*api.Service, error) {
    95  	conn, err := grpc.Dial(b.cfg.Manager, grpc.WithInsecure())
    96  	if err != nil {
    97  		return nil, err
    98  	}
    99  	client := api.NewControlClient(conn)
   100  	r, err := client.CreateService(ctx, &api.CreateServiceRequest{
   101  		Spec: b.spec(),
   102  	})
   103  	if err != nil {
   104  		return nil, err
   105  	}
   106  	return r.Service, nil
   107  }