github.com/artyom/thrift@v0.0.0-20130902103359-388840a05deb/socket.go (about) 1 /* 2 * Licensed to the Apache Software Foundation (ASF) under one 3 * or more contributor license agreements. See the NOTICE file 4 * distributed with this work for additional information 5 * regarding copyright ownership. The ASF licenses this file 6 * to you under the Apache License, Version 2.0 (the 7 * "License"); you may not use this file except in compliance 8 * with the License. You may obtain a copy of the License at 9 * 10 * http://www.apache.org/licenses/LICENSE-2.0 11 * 12 * Unless required by applicable law or agreed to in writing, 13 * software distributed under the License is distributed on an 14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 15 * KIND, either express or implied. See the License for the 16 * specific language governing permissions and limitations 17 * under the License. 18 */ 19 20 package thrift 21 22 import ( 23 "net" 24 "time" 25 ) 26 27 type TSocket struct { 28 conn net.Conn 29 addr net.Addr 30 timeout time.Duration 31 } 32 33 // NewTSocket creates a net.Conn-backed TTransport, given a host and port 34 // 35 // Example: 36 // trans, err := thrift.NewTSocket("localhost:9090") 37 func NewTSocket(hostPort string) (*TSocket, error) { 38 return NewTSocketTimeout(hostPort, 0) 39 } 40 41 // NewTSocketTimeout creates a net.Conn-backed TTransport, given a host and port 42 // it also accepts a timeout as a time.Duration 43 func NewTSocketTimeout(hostPort string, timeout time.Duration) (*TSocket, error) { 44 //conn, err := net.DialTimeout(network, address, timeout) 45 addr, err := net.ResolveTCPAddr("tcp", hostPort) 46 if err != nil { 47 return nil, err 48 } 49 return NewTSocketFromAddrTimeout(addr, timeout), nil 50 } 51 52 // Creates a TSocket from a net.Addr 53 func NewTSocketFromAddrTimeout(addr net.Addr, timeout time.Duration) *TSocket { 54 return &TSocket{addr: addr, timeout: timeout} 55 } 56 57 // Creates a TSocket from an existing net.Conn 58 func NewTSocketFromConnTimeout(conn net.Conn, timeout time.Duration) *TSocket { 59 return &TSocket{conn: conn, addr: conn.RemoteAddr(), timeout: timeout} 60 } 61 62 // Sets the socket timeout 63 func (p *TSocket) SetTimeout(timeout time.Duration) error { 64 p.timeout = timeout 65 return nil 66 } 67 68 func (p *TSocket) pushDeadline(read, write bool) { 69 var t time.Time 70 if p.timeout > 0 { 71 t = time.Now().Add(time.Duration(p.timeout)) 72 } 73 if read && write { 74 p.conn.SetDeadline(t) 75 } else if read { 76 p.conn.SetReadDeadline(t) 77 } else if write { 78 p.conn.SetWriteDeadline(t) 79 } 80 } 81 82 // Connects the socket, creating a new socket object if necessary. 83 func (p *TSocket) Open() error { 84 if p.IsOpen() { 85 return NewTTransportException(ALREADY_OPEN, "Socket already connected.") 86 } 87 if p.addr == nil { 88 return NewTTransportException(NOT_OPEN, "Cannot open nil address.") 89 } 90 if len(p.addr.Network()) == 0 { 91 return NewTTransportException(NOT_OPEN, "Cannot open bad network name.") 92 } 93 if len(p.addr.String()) == 0 { 94 return NewTTransportException(NOT_OPEN, "Cannot open bad address.") 95 } 96 var err error 97 if p.conn, err = net.DialTimeout(p.addr.Network(), p.addr.String(), p.timeout); err != nil { 98 return NewTTransportException(NOT_OPEN, err.Error()) 99 } 100 return nil 101 } 102 103 // Retreive the underlying net.Conn 104 func (p *TSocket) Conn() net.Conn { 105 return p.conn 106 } 107 108 // Returns true if the connection is open 109 func (p *TSocket) IsOpen() bool { 110 if p.conn == nil { 111 return false 112 } 113 return true 114 } 115 116 // Closes the socket. 117 func (p *TSocket) Close() error { 118 // Close the socket 119 if p.conn != nil { 120 err := p.conn.Close() 121 if err != nil { 122 return err 123 } 124 p.conn = nil 125 } 126 return nil 127 } 128 129 func (p *TSocket) Read(buf []byte) (int, error) { 130 if !p.IsOpen() { 131 return 0, NewTTransportException(NOT_OPEN, "Connection not open") 132 } 133 p.pushDeadline(true, false) 134 n, err := p.conn.Read(buf) 135 return n, NewTTransportExceptionFromError(err) 136 } 137 138 func (p *TSocket) Write(buf []byte) (int, error) { 139 if !p.IsOpen() { 140 return 0, NewTTransportException(NOT_OPEN, "Connection not open") 141 } 142 p.pushDeadline(false, true) 143 return p.conn.Write(buf) 144 } 145 146 func (p *TSocket) Peek() bool { 147 return p.IsOpen() 148 } 149 150 func (p *TSocket) Flush() error { 151 return nil 152 } 153 154 func (p *TSocket) Interrupt() error { 155 if !p.IsOpen() { 156 return nil 157 } 158 return p.conn.Close() 159 }