github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/ngaut/go-zookeeper/zk/tracer.go (about)

     1  package zk
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"sync"
     9  )
    10  
    11  var (
    12  	requests     = make(map[int32]int32) // Map of Xid -> Opcode
    13  	requestsLock = &sync.Mutex{}
    14  )
    15  
    16  func trace(conn1, conn2 net.Conn, client bool) {
    17  	defer conn1.Close()
    18  	defer conn2.Close()
    19  	buf := make([]byte, 10*1024)
    20  	init := true
    21  	for {
    22  		_, err := io.ReadFull(conn1, buf[:4])
    23  		if err != nil {
    24  			fmt.Println("1>", client, err)
    25  			return
    26  		}
    27  
    28  		blen := int(binary.BigEndian.Uint32(buf[:4]))
    29  
    30  		_, err = io.ReadFull(conn1, buf[4:4+blen])
    31  		if err != nil {
    32  			fmt.Println("2>", client, err)
    33  			return
    34  		}
    35  
    36  		var cr interface{}
    37  		opcode := int32(-1)
    38  		readHeader := true
    39  		if client {
    40  			if init {
    41  				cr = &connectRequest{}
    42  				readHeader = false
    43  			} else {
    44  				xid := int32(binary.BigEndian.Uint32(buf[4:8]))
    45  				opcode = int32(binary.BigEndian.Uint32(buf[8:12]))
    46  				requestsLock.Lock()
    47  				requests[xid] = opcode
    48  				requestsLock.Unlock()
    49  				cr = requestStructForOp(opcode)
    50  				if cr == nil {
    51  					fmt.Printf("Unknown opcode %d\n", opcode)
    52  				}
    53  			}
    54  		} else {
    55  			if init {
    56  				cr = &connectResponse{}
    57  				readHeader = false
    58  			} else {
    59  				xid := int32(binary.BigEndian.Uint32(buf[4:8]))
    60  				zxid := int64(binary.BigEndian.Uint64(buf[8:16]))
    61  				errnum := int32(binary.BigEndian.Uint32(buf[16:20]))
    62  				if xid != -1 || zxid != -1 {
    63  					requestsLock.Lock()
    64  					found := false
    65  					opcode, found = requests[xid]
    66  					if !found {
    67  						println("WEFWEFEW")
    68  						opcode = 0
    69  					}
    70  					delete(requests, xid)
    71  					requestsLock.Unlock()
    72  				} else {
    73  					opcode = opWatcherEvent
    74  				}
    75  				cr = responseStructForOp(opcode)
    76  				if cr == nil {
    77  					fmt.Printf("Unknown opcode %d\n", opcode)
    78  				}
    79  				if errnum != 0 {
    80  					cr = &struct{}{}
    81  				}
    82  			}
    83  		}
    84  		opname := "."
    85  		if opcode != -1 {
    86  			opname = opNames[opcode]
    87  		}
    88  		if cr == nil {
    89  			fmt.Printf("%+v %s %+v\n", client, opname, buf[4:4+blen])
    90  		} else {
    91  			n := 4
    92  			hdrStr := ""
    93  			if readHeader {
    94  				var hdr interface{}
    95  				if client {
    96  					hdr = &requestHeader{}
    97  				} else {
    98  					hdr = &responseHeader{}
    99  				}
   100  				if n2, err := decodePacket(buf[n:n+blen], hdr); err != nil {
   101  					fmt.Println(err)
   102  				} else {
   103  					n += n2
   104  				}
   105  				hdrStr = fmt.Sprintf(" %+v", hdr)
   106  			}
   107  			if _, err := decodePacket(buf[n:n+blen], cr); err != nil {
   108  				fmt.Println(err)
   109  			}
   110  			fmt.Printf("%+v %s%s %+v\n", client, opname, hdrStr, cr)
   111  		}
   112  
   113  		init = false
   114  
   115  		written, err := conn2.Write(buf[:4+blen])
   116  		if err != nil {
   117  			fmt.Println("3>", client, err)
   118  			return
   119  		} else if written != 4+blen {
   120  			fmt.Printf("Written != read: %d != %d\n", written, blen)
   121  			return
   122  		}
   123  	}
   124  }
   125  
   126  func handleConnection(addr string, conn net.Conn) {
   127  	zkConn, err := net.Dial("tcp", addr)
   128  	if err != nil {
   129  		fmt.Println(err)
   130  		return
   131  	}
   132  	go trace(conn, zkConn, true)
   133  	trace(zkConn, conn, false)
   134  }
   135  
   136  func StartTracer(listenAddr, serverAddr string) {
   137  	ln, err := net.Listen("tcp", listenAddr)
   138  	if err != nil {
   139  		panic(err)
   140  	}
   141  	for {
   142  		conn, err := ln.Accept()
   143  		if err != nil {
   144  			fmt.Println(err)
   145  			continue
   146  		}
   147  		go handleConnection(serverAddr, conn)
   148  	}
   149  }