add Set ADT and TCP Server

Reinvented some wheels. These should be useful later on.

Author: tabjy

LICENSE.md LICENSE

@@ -1,7 +1,6 @@
+MIT License
 
-The MIT License (MIT)
-
-Copyright (c) 2017 Tab/jy
+Copyright (c) 2017 Tabjy
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -19,4 +18,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+SOFTWARE.

common/set.go

@@ -0,0 +1,109 @@
+package common
+
+import "sync"
+
+// Set is a ADT that can store certain values, without any particular order,
+// and no repeated values. It can be understood as a unordered list with all
+// elements being distinct.
+type Set interface {
+	Add(element interface{})           // adds the specified element to this set if it is not already present.
+	Remove(element interface{}) bool   // removes the specified element from this set if it is present
+	Contains(element interface{}) bool // returns true if this set contains the specified element.
+	Clear()                            // removes all of the elements from this set
+	Len() int                          // returns the number of elements in this set (its cardinality).
+	Empty() bool                       // returns true if this set contains no elements.
+
+	ForEach(fn func(interface{}))     // apply callback function to each of the elements within the set.
+	Filter(fn func(interface{}) bool) // apply callback function to each of the elements within the set, a element is kept if and only if callback returns true.
+}
+
+// HashSet is a general implement of Set ADT with a underlying golang built-in
+// map. By natures of golang maps (where as order of iteration is undefined),
+// iterators are not possible to implement.
+type HashSet struct {
+	items map[interface{}]bool
+	mu    sync.Mutex
+}
+
+func NewHashSet() Set {
+	return &HashSet{items: make(map[interface{}]bool)}
+}
+
+// Add implements Add in Set ADT.
+func (s *HashSet) Add(element interface{}) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.items[element] = true
+}
+
+// Remove implements Remove in Set ADT.
+func (s *HashSet) Remove(element interface{}) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	_, contains := s.items[element]
+
+	if contains {
+		delete(s.items, element)
+	}
+	return contains
+}
+
+// Contains implements Contains in Set ADT.
+func (s *HashSet) Contains(element interface{}) bool {
+	// not sure is map lookup is thread safe, but just better be safe than sorry
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	_, contains := s.items[element]
+	return contains
+}
+
+// Clear implements Clear in Set ADT.
+func (s *HashSet) Clear() {
+	// the old map should be recycled by GC, hopefully...
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.items = make(map[interface{}]bool)
+}
+
+// Len implements Len in Set ADT.
+func (s *HashSet) Len() int {
+	// this is atomic already
+	return len(s.items)
+}
+
+// Empty implements Empty in Set ADT.
+func (s *HashSet) Empty() bool {
+	return s.Len() == 0
+}
+
+// ForEach implements ForEach in Set ADT.
+//
+// IMPORTANT: ForEach causes the set to be locked until finish iterating all
+// elements. Therefore, calling any of Add, Remove, Contains, and Clear will
+// result in a dead lock. Do these ops in a separate goroutine!
+func (s *HashSet) ForEach(fn func(interface{})) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for element := range s.items {
+		fn(element)
+	}
+}
+
+// Filter implements Filter in Set ADT.
+//
+// IMPORTANT: Filter causes the set to be locked until finish iterating all
+// elements. Therefore, calling any of Add, Remove, Contains, and Clear will
+// result in a dead lock. Do these ops in a separate goroutine!
+func (s *HashSet) Filter(fn func(interface{}) bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for element := range s.items {
+		if keep := fn(element); !keep {
+			delete(s.items, element)
+		}
+	}
+}

common/tcp.go

@@ -0,0 +1,236 @@
+// Package common provides a few generalized interface/struct utilities might
+// be handy for other packages. The goal is to reduce code redundancy.
+package common
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+
+	"github.com/tabjy/yagl"
+)
+
+// ErrServerClosed is returned by the TCPServer's Serve, and ListenAndServe,
+// methods after a call to Shutdown or Close.
+var ErrServerClosed = errors.New("common: TCPServer closed")
+
+// ErrServerNotListening is returned by TCPServer's Serve, and ListenAndServe,
+// methods if being called before calling Listen.
+var ErrServerNotListening = errors.New("common: TCPServer not listening")
+
+type TCPHandler interface {
+	ServeTCP(ctx context.Context, conn net.Conn)
+}
+
+type echoHandler struct{}
+
+func (h *echoHandler) ServeTCP(ctx context.Context, conn net.Conn) {
+	go func() {
+		<-ctx.Done() // this doesn't block forever, TCPServer call cancel after ServeTCP returns
+		// conn.Close() // could fail with error, but it's okay
+		yagl.Trace("echoHandler goroutine unblocks and exists")
+	}()
+	if _, err := io.Copy(conn, conn); err != nil {
+		if ctx.Err() != nil {
+			yagl.Infof("context canceled: %v", ctx.Err())
+		} else {
+			yagl.Errorf("fail to echo request: %v", err)
+		}
+	}
+}
+
+var EchoHandler = &echoHandler{}
+
+// A TCPServer defines parameters for running a TCP server. The zero value for
+// TCPServer is a valid configuration.
+type TCPServer struct {
+	Host string // IP address or hostname to listen on. Leave empty for an unspecified address.
+	Port uint16 // Port to listen on. A port number is automatically chosen if left empty or 0.
+
+	Handler TCPHandler
+
+	// ErrorLog specifies an optional logger
+	// If nil, logging goes to os.Stderr via a yagl standard logger
+	Logger yagl.Logger
+
+	ln    net.Listener
+	conns Set
+
+	ctx    context.Context
+	cancel context.CancelFunc
+	wg     sync.WaitGroup
+}
+
+func (srv *TCPServer) logger() yagl.Logger {
+	if srv.Logger == nil {
+		return yagl.StdLogger()
+	}
+	return srv.Logger
+}
+
+// Listen listens on srv.Host:srv.Port. If a Listener is ready created, the old
+// one will be closed and replaced.
+func (srv *TCPServer) Listen() error {
+	addr := net.JoinHostPort(srv.Host, strconv.Itoa(int(srv.Port)))
+
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		srv.logger().Errorf("failed to listen on %s: %v", addr, err)
+		return err
+	}
+	srv.ln = ln
+	srv.logger().Infof("TCPServer listening on %v", srv.ln.Addr())
+
+	return nil
+}
+
+// Serve accepts incoming connections on the Listener ln, creating a new
+// service goroutine for each. The service goroutines read requests and then
+// call srv.Handler to handle to them. Make sure Listen is called before
+// calling this function.
+//
+// Serve always returns a non-nil error. After Shutdown or Close, the returned
+// error is ErrServerClosed.
+func (srv *TCPServer) Serve() error {
+	if srv.ln == nil {
+		return ErrServerNotListening
+	}
+
+	srv.conns = NewHashSet()
+
+	if srv.Handler == nil {
+		srv.Handler = EchoHandler
+	}
+
+	srv.ctx, srv.cancel = context.WithCancel(context.Background())
+
+	for true {
+		conn, err := srv.ln.Accept()
+		if err != nil {
+			// server level error, causing server to stop
+			// ln.Accept unblocks and returns error when ln.Close called, by design
+			if strings.Contains(err.Error(), "use of closed network connection") {
+
+				return ErrServerClosed
+			}
+			srv.logger().Errorf("TCPServer stopping for error: %v", err)
+			return err
+		}
+
+		srv.wg.Add(1)
+		go func() {
+			defer func() {
+				srv.wg.Add(-1)
+				srv.conns.Remove(conn)
+			}()
+			srv.conns.Add(conn)
+			srv.logger().Tracef("new connection from %v", conn.RemoteAddr())
+
+			srv.logger().Tracef("connection to be handled by %T", srv.Handler)
+			ctx, cancel := context.WithCancel(srv.ctx)
+			srv.Handler.ServeTCP(ctx, conn)
+			srv.logger().Tracef("handler %T returned, connection closing...", srv.Handler)
+			cancel()
+
+			// try to closed connection, even if it's closed by handler already
+			if err := conn.Close(); err != nil {
+				// test if because conn already closed
+				if !strings.Contains(err.Error(), "use of closed network connection") {
+					srv.logger().Panicf("failed to close connection from %v, %v", conn.RemoteAddr(), err)
+					// goroutine stops here
+				}
+				srv.logger().Warnf("connection from %v is already closed", conn.RemoteAddr())
+			}
+			srv.logger().Tracef("connection from %v is now closed", conn.RemoteAddr())
+
+
+		}()
+	}
+
+	// this never happens due to the infinite loop
+	return nil
+}
+
+// ListenAndServe first call Listen, then calls Server to handle incoming
+// connections. If srv.Addr is blank, ":tcp" is used.
+//
+// ListenAndServe always returns a non-nil error. After Shutdown or Close, the
+// returned error is ErrServerClosed.
+func (srv *TCPServer) ListenAndServe() error {
+	if err := srv.Listen(); err != nil {
+		return err
+	}
+
+	return srv.Serve()
+}
+
+func (srv *TCPServer) forceCloseConns() {
+	srv.logger().Tracef("forcing to close all connections, %d remaining", srv.conns.Len())
+	srv.conns.ForEach(func(element interface{}) {
+		conn, _ := element.(net.Conn)
+		srv.logger().Tracef("forcing to close %v", conn.RemoteAddr())
+		if err := conn.Close(); err != nil {
+			// connection level errors, no need to deal with them, just log
+			srv.logger().Errorf("failed to close connection: %v", err)
+		}
+	})
+	srv.conns.Clear()
+	srv.wg.Done()
+}
+
+// Close immediately closes active net.Listener and closes all active
+// connections. This could be unsafe as there might be ongoing goroutines
+// handling connections. Close return after Serve returns.
+//
+// Close returns any error returned from closing the Server's underlying
+// Listener(s).
+func (srv *TCPServer) Close() error {
+	srv.logger().Infof("closing server listening on %v", srv.ln.Addr())
+
+	// first close listener, so no more incoming connections
+	if err := srv.ln.Close(); err != nil {
+		srv.logger().Errorf("failed to close listener: %v", err)
+		return err
+	}
+
+	srv.cancel() // notify all handlers to finish whatever is left
+	srv.forceCloseConns()
+
+	return nil
+}
+
+// Shutdown immediately closes active net.Listener and sends close signals to
+// all actives connection through context; still, it's up to connections'
+// handlers to decide what to do with close signals. Shutdown wait for all
+// connection to be closed before returning.
+//
+// Best practice: wait for a certain amount of time, then just call Close to
+// force close all connections.
+//
+// Shutdown returns any error returned from closing the Server's underlying
+// Listener(s).
+func (srv *TCPServer) Shutdown() error {
+	srv.logger().Infof("shutting down server listening on %v", srv.ln.Addr())
+
+	// first close listener, so no more incoming connections
+	if err := srv.ln.Close(); err != nil {
+		srv.logger().Errorf("failed to close listener: %v", err)
+		return err
+	}
+
+	srv.cancel() // notify all handler to finish whatever is left
+	srv.logger().Infof("shutdown signal sent, waiting for all connection to be closed, %d remaining", srv.conns.Len())
+	srv.wg.Wait()
+	// by the time wg.Wait unblocks, all connection SHOULD be closed
+	// but let's just check for sure, for debug purpose
+	if srv.conns.Len() != 0 {
+		srv.Logger.Warn("not all connection are closed upon Shutdown!")
+		srv.forceCloseConns()
+	}
+
+	return nil
+}