go.dedis.ch/onet/v4@v4.0.0-pre1/simul/monitor/tcpproxy.go (about)

     1  // Copyright 2016 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package monitor
    16  
    17  import (
    18  	"fmt"
    19  	"io"
    20  	"math/rand"
    21  	"net"
    22  	"sync"
    23  	"time"
    24  
    25  	"go.dedis.ch/onet/v4/log"
    26  	"golang.org/x/xerrors"
    27  )
    28  
    29  type remote struct {
    30  	mu       sync.Mutex
    31  	srv      *net.SRV
    32  	addr     string
    33  	inactive bool
    34  }
    35  
    36  func (r *remote) inactivate() {
    37  	r.mu.Lock()
    38  	defer r.mu.Unlock()
    39  	r.inactive = true
    40  }
    41  
    42  func (r *remote) tryReactivate() error {
    43  	conn, err := net.Dial("tcp", r.addr)
    44  	if err != nil {
    45  		return xerrors.Errorf("dial: %v", err)
    46  	}
    47  	conn.Close()
    48  	r.mu.Lock()
    49  	defer r.mu.Unlock()
    50  	r.inactive = false
    51  	return nil
    52  }
    53  
    54  func (r *remote) isActive() bool {
    55  	r.mu.Lock()
    56  	defer r.mu.Unlock()
    57  	return !r.inactive
    58  }
    59  
    60  // A TCPProxy proxies connections arriving at the Listener to one of
    61  // the Endpoints.
    62  type TCPProxy struct {
    63  	Listener        net.Listener
    64  	Endpoints       []*net.SRV
    65  	MonitorInterval time.Duration
    66  
    67  	donec chan struct{}
    68  
    69  	mu        sync.Mutex // guards the following fields
    70  	remotes   []*remote
    71  	pickCount int // for round robin
    72  }
    73  
    74  // Run starts servicing clients. It will not return until Stop is called.
    75  func (tp *TCPProxy) Run() error {
    76  	tp.donec = make(chan struct{})
    77  	if tp.MonitorInterval == 0 {
    78  		tp.MonitorInterval = 5 * time.Minute
    79  	}
    80  	for _, srv := range tp.Endpoints {
    81  		addr := fmt.Sprintf("%s:%d", srv.Target, srv.Port)
    82  		tp.remotes = append(tp.remotes, &remote{srv: srv, addr: addr})
    83  	}
    84  
    85  	eps := []string{}
    86  	for _, ep := range tp.Endpoints {
    87  		eps = append(eps, fmt.Sprintf("%s:%d", ep.Target, ep.Port))
    88  	}
    89  
    90  	go tp.runMonitor()
    91  	for {
    92  		in, err := tp.Listener.Accept()
    93  		if err != nil {
    94  			return xerrors.Errorf("listening: %v", err)
    95  		}
    96  
    97  		go tp.serve(in)
    98  	}
    99  }
   100  
   101  func (tp *TCPProxy) pick() *remote {
   102  	var weighted []*remote
   103  	var unweighted []*remote
   104  
   105  	bestPr := uint16(65535)
   106  	w := 0
   107  	// find best priority class
   108  	for _, r := range tp.remotes {
   109  		switch {
   110  		case !r.isActive():
   111  		case r.srv.Priority < bestPr:
   112  			bestPr = r.srv.Priority
   113  			w = 0
   114  			weighted = nil
   115  			unweighted = []*remote{r}
   116  			fallthrough
   117  		case r.srv.Priority == bestPr:
   118  			if r.srv.Weight > 0 {
   119  				weighted = append(weighted, r)
   120  				w += int(r.srv.Weight)
   121  			} else {
   122  				unweighted = append(unweighted, r)
   123  			}
   124  		}
   125  	}
   126  	if weighted != nil {
   127  		if len(unweighted) > 0 && rand.Intn(100) == 1 {
   128  			// In the presence of records containing weights greater
   129  			// than 0, records with weight 0 should have a very small
   130  			// chance of being selected.
   131  			r := unweighted[tp.pickCount%len(unweighted)]
   132  			tp.pickCount++
   133  			return r
   134  		}
   135  		// choose a uniform random number between 0 and the sum computed
   136  		// (inclusive), and select the RR whose running sum value is the
   137  		// first in the selected order
   138  		choose := rand.Intn(w)
   139  		for i := 0; i < len(weighted); i++ {
   140  			choose -= int(weighted[i].srv.Weight)
   141  			if choose <= 0 {
   142  				return weighted[i]
   143  			}
   144  		}
   145  	}
   146  	if unweighted != nil {
   147  		for i := 0; i < len(tp.remotes); i++ {
   148  			picked := tp.remotes[tp.pickCount%len(tp.remotes)]
   149  			tp.pickCount++
   150  			if picked.isActive() {
   151  				return picked
   152  			}
   153  		}
   154  	}
   155  	return nil
   156  }
   157  
   158  func (tp *TCPProxy) serve(in net.Conn) {
   159  	var (
   160  		err error
   161  		out net.Conn
   162  	)
   163  
   164  	for {
   165  		tp.mu.Lock()
   166  		remote := tp.pick()
   167  		tp.mu.Unlock()
   168  		if remote == nil {
   169  			break
   170  		}
   171  		// TODO: add timeout
   172  		out, err = net.Dial("tcp", remote.addr)
   173  		if err == nil {
   174  			break
   175  		}
   176  		remote.inactivate()
   177  		log.Warnf("deactivated endpoint [%s] due to %v for %v", remote.addr, err, tp.MonitorInterval)
   178  	}
   179  
   180  	if out == nil {
   181  		in.Close()
   182  		return
   183  	}
   184  
   185  	go func() {
   186  		io.Copy(in, out)
   187  		in.Close()
   188  		out.Close()
   189  	}()
   190  
   191  	io.Copy(out, in)
   192  	out.Close()
   193  	in.Close()
   194  }
   195  
   196  func (tp *TCPProxy) runMonitor() {
   197  	for {
   198  		select {
   199  		case <-time.After(tp.MonitorInterval):
   200  			tp.mu.Lock()
   201  			for _, rem := range tp.remotes {
   202  				if rem.isActive() {
   203  					continue
   204  				}
   205  				go func(r *remote) {
   206  					if err := r.tryReactivate(); err != nil {
   207  						log.Warnf("failed to activate endpoint [%s] due to %v (stay inactive for another %v)", r.addr, err, tp.MonitorInterval)
   208  					} else {
   209  						log.Infof("activated %s", r.addr)
   210  					}
   211  				}(rem)
   212  			}
   213  			tp.mu.Unlock()
   214  		case <-tp.donec:
   215  			return
   216  		}
   217  	}
   218  }
   219  
   220  // Stop stops a running TCPProxy. After calling Stop, Run will return.
   221  func (tp *TCPProxy) Stop() {
   222  	// graceful shutdown?
   223  	// shutdown current connections?
   224  	tp.Listener.Close()
   225  	close(tp.donec)
   226  }