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 }