github.com/mailgun/holster/v4@v4.20.0/election/example_test.go (about)

     1  package election_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"log"
     9  	"net/http"
    10  	"os"
    11  	"os/signal"
    12  	"syscall"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/mailgun/holster/v4/election"
    17  	"github.com/mailgun/holster/v4/errors"
    18  	"github.com/sirupsen/logrus"
    19  	"github.com/stretchr/testify/require"
    20  )
    21  
    22  func sendRPC(ctx context.Context, peer string, req election.RPCRequest, resp *election.RPCResponse) error {
    23  	// Marshall the RPC request to json
    24  	b, err := json.Marshal(req)
    25  	if err != nil {
    26  		return errors.Wrap(err, "while encoding request")
    27  	}
    28  
    29  	// Create a new http request with context
    30  	hr, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://%s/rpc", peer), bytes.NewBuffer(b))
    31  	if err != nil {
    32  		return errors.Wrap(err, "while creating request")
    33  	}
    34  
    35  	// Send the request
    36  	hp, err := http.DefaultClient.Do(hr)
    37  	if err != nil {
    38  		return errors.Wrap(err, "while sending http request")
    39  	}
    40  	defer func() {
    41  		_ = hp.Body.Close()
    42  	}()
    43  
    44  	// Decode the response from JSON
    45  	dec := json.NewDecoder(hp.Body)
    46  	if err := dec.Decode(&resp); err != nil {
    47  		return errors.Wrap(err, "while decoding response")
    48  	}
    49  	return nil
    50  }
    51  
    52  func newHandler(t *testing.T, node election.Node) func(w http.ResponseWriter, r *http.Request) {
    53  	return func(w http.ResponseWriter, r *http.Request) {
    54  		dec := json.NewDecoder(r.Body)
    55  		var req election.RPCRequest
    56  		if err := dec.Decode(&req); err != nil {
    57  			w.WriteHeader(http.StatusBadRequest)
    58  			_, err = w.Write([]byte(err.Error()))
    59  			require.NoError(t, err)
    60  		}
    61  		var resp election.RPCResponse
    62  		node.ReceiveRPC(req, &resp)
    63  
    64  		enc := json.NewEncoder(w)
    65  		if err := enc.Encode(resp); err != nil {
    66  			w.WriteHeader(http.StatusInternalServerError)
    67  			_, err = w.Write([]byte(err.Error()))
    68  			require.NoError(t, err)
    69  		}
    70  	}
    71  }
    72  
    73  // This example spawns 2 nodes, in a real application you would
    74  // only spawn a single node which would represent your application
    75  // in the election.
    76  func SimpleExample(t *testing.T) {
    77  	logrus.SetLevel(logrus.DebugLevel)
    78  
    79  	node1, err := election.NewNode(election.Config{
    80  		// A list of known peers at startup
    81  		Peers: []string{"localhost:7080", "localhost:7081"},
    82  		// A unique identifier used to identify us in a list of peers
    83  		UniqueID: "localhost:7080",
    84  		// Called whenever the library detects a change in leadership
    85  		OnUpdate: func(leader string) {
    86  			log.Printf("Current Leader: %s\n", leader)
    87  		},
    88  		// Called when the library wants to contact other peers
    89  		SendRPC: sendRPC,
    90  	})
    91  	if err != nil {
    92  		log.Fatal(err)
    93  	}
    94  	defer func() {
    95  		err := node1.Stop(context.Background())
    96  		require.NoError(t, err)
    97  	}()
    98  
    99  	node2, err := election.NewNode(election.Config{
   100  		Peers:    []string{"localhost:7080", "localhost:7081"},
   101  		UniqueID: "localhost:7081",
   102  		SendRPC:  sendRPC,
   103  	})
   104  	if err != nil {
   105  		t.Fatal(err)
   106  	}
   107  
   108  	go func() {
   109  		mux := http.NewServeMux()
   110  		mux.HandleFunc("/rpc", newHandler(t, node1))
   111  		server := &http.Server{
   112  			Addr:              ":7080",
   113  			Handler:           mux,
   114  			ReadHeaderTimeout: 1 * time.Minute,
   115  		}
   116  		log.Fatal(server.ListenAndServe())
   117  	}()
   118  
   119  	go func() {
   120  		mux := http.NewServeMux()
   121  		mux.HandleFunc("/rpc", newHandler(t, node2))
   122  		server := &http.Server{
   123  			Addr:              ":7081",
   124  			Handler:           mux,
   125  			ReadHeaderTimeout: 1 * time.Minute,
   126  		}
   127  		log.Fatal(server.ListenAndServe())
   128  	}()
   129  
   130  	// Wait for each of the http listeners to start fielding requests
   131  	if err := election.WaitForConnect("localhost:7080", 3, time.Second); err != nil {
   132  		t.Fatal(err)
   133  	}
   134  
   135  	if err := election.WaitForConnect("localhost:7081", 3, time.Second); err != nil {
   136  		t.Fatal(err)
   137  	}
   138  
   139  	// Now that both http handlers are listening for requests we
   140  	// can safely start the election.
   141  	err = node1.Start(context.Background())
   142  	require.NoError(t, err)
   143  	err = node2.Start(context.Background())
   144  	require.NoError(t, err)
   145  
   146  	// Wait here for signals to clean up our mess
   147  	c := make(chan os.Signal, 1)
   148  	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
   149  	for range c {
   150  		err = node1.Stop(context.Background())
   151  		require.NoError(t, err)
   152  		err = node2.Stop(context.Background())
   153  		require.NoError(t, err)
   154  		break
   155  	}
   156  }