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  }