package pingo

import (
    "bufio"
    "errors"
    "fmt"
    "io"
    "log"
    "net"
    "net/rpc"
    "os"
    "os/exec"
    "strings"
    "time"
)

var (
    errInvalidMessage      = ErrInvalidMessage(errors.New("Invalid ready message"))
    errRegistrationTimeout = ErrRegistrationTimeout(errors.New("Registration timed out"))
)

// Represents a plugin. After being created the plugin is not started or ready to run.
//
// Additional configuration (ErrorHandler and Timeout) can be set after initialization.
//
// Use Start() to make the plugin available.
type Plugin struct {
    exe         string
    proto       string
    unixdir     string
    params      []string
    initTimeout time.Duration
    exitTimeout time.Duration
    handler     ErrorHandler
    running     bool
    meta        meta
    objsCh      chan *objects
    connCh      chan *conn
    killCh      chan *waiter
    exitCh      chan struct{}
}

// NewPlugin create a new plugin ready to be started, or returns an error if the initial setup fails.
//
// The first argument specifies the protocol. It can be either set to "unix" for communication on an
// ephemeral local socket, or "tcp" for network communication on the local host (using a random
// unprivileged port.)
//
// This constructor will panic if the proto argument is neither "unix" nor "tcp".
//
// The path to the plugin executable should be absolute. Any path accepted by the "exec" package in the
// standard library is accepted and the same rules for execution are applied.
//
// Optionally some parameters might be passed to the plugin executable.
func NewPlugin(proto, path string, params ...string) *Plugin {
    if proto != "unix" && proto != "tcp" {
        panic("Invalid protocol. Specify 'unix' or 'tcp'.")
    }
    p := &Plugin{
        exe:         path,
        proto:       proto,
        params:      params,
        initTimeout: 2 * time.Second,
        exitTimeout: 2 * time.Second,
        handler:     NewDefaultErrorHandler(),
        meta:        meta("pingo" + randstr(5)),
        objsCh:      make(chan *objects),
        connCh:      make(chan *conn),
        killCh:      make(chan *waiter),
        exitCh:      make(chan struct{}),
    }
    return p
}

// Set the error (and output) handler implementation.  Use this to set a custom implementation.
// By default, standard logging is used.  See ErrorHandler.
//
// Panics if called after Start.
func (p *Plugin) SetErrorHandler(h ErrorHandler) {
    if p.running {
        panic("Cannot call SetErrorHandler after Start")
    }
    p.handler = h
}

// Set the maximum time a plugin is allowed to start up and to shut down.  Empty timeout (zero)
// is not allowed, default will be used.
//
// Default is two seconds.
//
// Panics if called after Start.
func (p *Plugin) SetTimeout(t time.Duration) {
    if p.running {
        panic("Cannot call SetTimeout after Start")
    }
    if t == 0 {
        return
    }
    p.initTimeout = t
    p.exitTimeout = t
}

func (p *Plugin) SetSocketDirectory(dir string) {
    if p.running {
        panic("Cannot call SetSocketDirectory after Start")
    }
    p.unixdir = dir
}

// Default string representation
func (p *Plugin) String() string {
    return fmt.Sprintf("%s %s", p.exe, strings.Join(p.params, " "))
}

// Start will execute the plugin as a subprocess. Start will return immediately. Any first call to the
// plugin will reveal eventual errors occurred at initialization.
//
// Calls subsequent to Start will hang until the plugin has been properly initialized.
func (p *Plugin) Start() {
    p.running = true
    go p.run()
}

// Stop attemps to stop cleanly or kill the running plugin, then will free all resources.
// Stop returns when the plugin as been shut down and related routines have exited.
func (p *Plugin) Stop() {
    wr := newWaiter()
    p.killCh <- wr
    wr.wait()
    p.exitCh <- struct{}{}
}

// Call performs an RPC call to the plugin. Prior to calling Call, the plugin must have been
// initialized by calling Start.
//
// Call will hang until a plugin has been initialized; it will return any error that happens
// either when performing the call or during plugin initialization via Start.
//
// Please refer to the "rpc" package from the standard library for more information on the
// semantics of this function.
func (p *Plugin) Call(name string, args interface{}, resp interface{}) error {
    conn := &conn{wr: newWaiter()}
    p.connCh <- conn
    conn.wr.wait()

    if conn.err != nil {
        return conn.err
    }

    return conn.client.Call(name, args, resp)
}

// Objects returns a list of the exported objects from the plugin. Exported objects used
// internally are not reported.
//
// Like Call, Objects returns any error happened on initialization if called after Start.
func (p *Plugin) Objects() ([]string, error) {
    objects := &objects{wr: newWaiter()}
    p.objsCh <- objects
    objects.wr.wait()

    return objects.list, objects.err
}

// ErrorHandler is the interface used by Plugin to report non-fatal errors and any other
// output from the plugin.
//
// A default implementation is provided and used if none is specified on plugin creation.
type ErrorHandler interface {
    // Error is called whenever a non-fatal error occurs in the plugin subprocess.
    Error(error)
    // Print is called for each line of output received from the plugin subprocess.
    Print(interface{})
}

// Default error handler implementation. Uses the default logging facility from the
// Go standard library.
type DefaultErrorHandler struct{}

// Constructor for default error handler.
func NewDefaultErrorHandler() *DefaultErrorHandler {
    return &DefaultErrorHandler{}
}

// Log via default standard library facility prepending the "error: " string.
func (e *DefaultErrorHandler) Error(err error) {
    log.Print("error: ", err)
}

// Log via default standard library facility.
func (e *DefaultErrorHandler) Print(s interface{}) {
    log.Print(s)
}

const internalObject = "PingoRpc"

type conn struct {
    client *rpc.Client
    err    error
    wr     *waiter
}

type waiter struct {
    c chan struct{}
}

func newWaiter() *waiter {
    return &waiter{c: make(chan struct{})}
}

func (wr *waiter) wait() {
    <-wr.c
}

func (wr *waiter) done() {
    close(wr.c)
}

func (wr *waiter) reset() {
    wr.c = make(chan struct{})
}

type client struct {
    *rpc.Client
    secret string
}

func newClient(s string, conn io.ReadWriteCloser) *client {
    return &client{secret: s, Client: rpc.NewClient(conn)}
}

func (c *client) authenticate(w io.Writer) error {
    _, err := io.WriteString(w, "Auth-Token: "+c.secret+"\n\n")
    return err
}

func dialAuthRpc(secret, network, address string, timeout time.Duration) (*rpc.Client, error) {
    conn, err := net.DialTimeout(network, address, timeout)
    if err != nil {
        return nil, err
    }
    c := newClient(secret, conn)
    if err := c.authenticate(conn); err != nil {
        return nil, err
    }
    return c.Client, nil
}

type objects struct {
    list []string
    err  error
    wr   *waiter
}

type ctrl struct {
    p    *Plugin
    objs []string
    // Protocol and address for RPC
    proto, addr string
    // Secret needed to connect to server
    secret string
    // Unrecoverable error is used as response to calls after it happened.
    err error
    // This channel is an alias to p.connCh. It allows to
    // intermittedly process calls (only when we can handle them).
    connCh chan *conn
    // Same as above, but for objects requests
    objsCh chan *objects
    // Timeout on plugin startup time
    timeoutCh <-chan time.Time
    // Get notification from Wait on the subprocess
    waitCh chan error
    // Get output lines from subprocess
    linesCh chan string
    // Respond to a routine waiting for this mail loop to exit.
    over *waiter
    // Executable
    proc *os.Process
    // RPC client to subprocess
    client *rpc.Client
}

func newCtrl(p *Plugin, t time.Duration) *ctrl {
    return &ctrl{
        p:         p,
        timeoutCh: time.After(t),
        linesCh:   make(chan string),
        waitCh:    make(chan error),
    }
}

func (c *ctrl) fatal(err error) {
    c.err = err
    c.open()
    c.kill()
}

func (c *ctrl) isFatal() bool {
    return c.err != nil
}

func (c *ctrl) close() {
    c.connCh = nil
    c.objsCh = nil
}

func (c *ctrl) open() {
    c.connCh = c.p.connCh
    c.objsCh = c.p.objsCh
}

func (c *ctrl) ready(val string) bool {
    var err error

    if err := c.parseReady(val); err != nil {
        c.fatal(err)
        return false
    }

    c.client, err = dialAuthRpc(c.secret, c.proto, c.addr, c.p.initTimeout)
    if err != nil {
        c.fatal(err)
        return false
    }

    // Remove the temp socket now that we are connected
    if c.proto == "unix" {
        if err := os.Remove(c.addr); err != nil {
            c.p.handler.Error(errors.New("Cannot remove temporary socket: " + err.Error()))
        }
    }

    // Defuse the timeout on ready
    c.timeoutCh = nil

    return true
}

func (c *ctrl) readOutput(r io.Reader) {
    scanner := bufio.NewScanner(r)

    for scanner.Scan() {
        c.linesCh <- scanner.Text()
    }
}

func (c *ctrl) waitErr(pidCh chan<- int, err error) {
    close(pidCh)
    c.waitCh <- err
}

func (c *ctrl) wait(pidCh chan<- int, exe string, params ...string) {
    defer close(c.waitCh)

    cmd := exec.Command(exe, params...)

    stdout, err := cmd.StdoutPipe()
    if err != nil {
        c.waitErr(pidCh, err)
        return
    }
    stderr, err := cmd.StderrPipe()
    if err != nil {
        c.waitErr(pidCh, err)
        return
    }
    if err := cmd.Start(); err != nil {
        c.waitErr(pidCh, err)
        return
    }

    pidCh <- cmd.Process.Pid
    close(pidCh)

    c.readOutput(stdout)
    c.readOutput(stderr)

    c.waitCh <- cmd.Wait()
}

func (c *ctrl) kill() {
    if c.proc == nil {
        return
    }
    // Ignore errors here because Kill might have been called after
    // process has ended.
    c.proc.Kill()
    c.proc = nil
}

func (c *ctrl) parseReady(str string) error {
    if !strings.HasPrefix(str, "proto=") {
        return errInvalidMessage
    }
    str = str[6:]
    s := strings.IndexByte(str, ' ')
    if s < 0 {
        return errInvalidMessage
    }
    proto := str[0:s]
    if proto != "unix" && proto != "tcp" {
        return errInvalidMessage
    }
    c.proto = proto

    str = str[s+1:]
    if !strings.HasPrefix(str, "addr=") {
        return errInvalidMessage
    }
    c.addr = str[5:]

    return nil
}

// Copy the list of objects for the requestor
func (c *ctrl) objects() []string {
    list := make([]string, len(c.objs)-1)
    for i, j := 0, 0; i < len(c.objs); i++ {
        if c.objs[i] == internalObject {
            continue
        }
        list[j] = c.objs[i]
        j = j + 1
    }
    return list
}

func (p *Plugin) run() {
    if p.unixdir == "" {
        p.unixdir = os.TempDir()
    }

    params := []string{
        "-pingo:prefix=" + string(p.meta),
        "-pingo:proto=" + p.proto,
    }
    if p.proto == "unix" && p.unixdir != "" {
        params = append(params, "-pingo:unixdir="+p.unixdir)
    }
    for i := 0; i < len(p.params); i++ {
        params = append(params, p.params[i])
    }

    c := newCtrl(p, p.initTimeout)

    pidCh := make(chan int)
    go c.wait(pidCh, p.exe, params...)
    pid := <-pidCh

    if pid != 0 {
        if proc, err := os.FindProcess(pid); err == nil {
            c.proc = proc
        }
    }

    for {
        select {
        case <-c.timeoutCh:
            c.fatal(errRegistrationTimeout)
        case r := <-c.connCh:
            if c.isFatal() {
                r.err = c.err
                r.wr.done()
                continue
            }

            r.client = c.client
            r.wr.done()
        case o := <-c.objsCh:
            if c.isFatal() {
                o.err = c.err
                o.wr.done()
                continue
            }

            o.list = c.objects()
            o.wr.done()
        case line := <-c.linesCh:
            key, val := p.meta.parse(line)
            switch key {
            case "auth-token":
                c.secret = val
            case "fatal":
                if err := parseError(val); err != nil {
                    c.fatal(err)
                } else {
                    c.fatal(errors.New(val))
                }
            case "error":
                if err := parseError(val); err != nil {
                    p.handler.Print(err)
                } else {
                    p.handler.Print(errors.New(val))
                }
            case "objects":
                c.objs = strings.Split(val, ", ")
            case "ready":
                if !c.ready(val) {
                    continue
                }
                // Start accepting calls
                c.open()
            default:
                p.handler.Print(line)
            }
        case wr := <-p.killCh:
            if c.waitCh == nil {
                wr.done()
                continue
            }

            // If we don't accept calls, kill immediately
            if c.connCh == nil || c.client == nil {
                c.kill()
            } else {
                // Be sure to kill the process if it doesn't obey Exit.
                go func(pid int, t time.Duration) {
                    <-time.After(t)

                    if proc, err := os.FindProcess(pid); err == nil {
                        proc.Kill()
                    }
                }(pid, p.exitTimeout)

                c.client.Call(internalObject+".Exit", 0, nil)
            }

            if c.client != nil {
                c.client.Close()
            }

            // Do not accept calls
            c.close()

            // When wait on the subprocess is exited, signal back via "over"
            c.over = wr
        case err := <-c.waitCh:
            if err != nil {
                if _, ok := err.(*exec.ExitError); !ok {
                    p.handler.Error(err)
                }
                c.fatal(err)
            }

            // Signal to whoever killed us (via killCh) that we are done
            if c.over != nil {
                c.over.done()
            }

            c.proc = nil
            c.waitCh = nil
            c.linesCh = nil
        case <-p.exitCh:
            return
        }
    }
}

plugin.go 源码阅读的更多相关文章

  1. 详细讲解Hadoop源码阅读工程（以hadoop-2.6.0-src.tar.gz和hadoop-2.6.0-cdh5.4.5-src.tar.gz为代表）

    首先,说的是,本人到现在为止,已经玩过.                   对于,这样的软件,博友,可以去看我博客的相关博文.在此,不一一赘述! Eclipse *版本 Eclipse *下载 Jd ...

  2. 内核源码阅读vim+cscope+ctags+taglist

    杜斌博客:http://blog.db89.org/kernel-source-read-vim-cscope-ctags-taglist/ 武特博客:http://edsionte.com/tech ...

  3. Caddy源码阅读（二）启动流程与 Event 事件通知

    Caddy源码阅读(二)启动流程与 Event 事件通知 Preface Caddy 是 Go 语言构建的轻量配置化服务器.https://github.com/caddyserver/caddy C ...

  4. react v16.12 源码阅读环境搭建

    搭建后的代码(Keep updated): https://github.com/lirongfei123/read-react 欢迎将源码阅读遇到的问题提到issue 环境搭建思路: 搭建一个web ...

  5. 搭建 Spring 源码阅读环境

    前言 有一个Spring源码阅读环境是学习Spring的基础.笔者借鉴了网上很多搭建环境的方法,也尝试了很多,接下来总结两种个人认为比较简便实用的方法.读者可根据自己的需要自行选择. 方法一:搭建基础 ...

  6. 【原】FMDB源码阅读（三）

    [原]FMDB源码阅读(三) 本文转载请注明出处 —— polobymulberry-博客园 1. 前言 FMDB比较优秀的地方就在于对多线程的处理.所以这一篇主要是研究FMDB的多线程处理的实现.而 ...

  7. 【原】FMDB源码阅读（二）

    [原]FMDB源码阅读(二) 本文转载请注明出处 -- polobymulberry-博客园 1. 前言 上一篇只是简单地过了一下FMDB一个简单例子的基本流程,并没有涉及到FMDB的所有方方面面,比 ...

  8. 【原】FMDB源码阅读（一）

    [原]FMDB源码阅读(一) 本文转载请注明出处 —— polobymulberry-博客园 1. 前言 说实话,之前的SDWebImage和AFNetworking这两个组件我还是使用过的,但是对于 ...

  9. 【原】AFNetworking源码阅读（六）

    [原]AFNetworking源码阅读(六) 本文转载请注明出处 —— polobymulberry-博客园 1. 前言 这一篇的想讲的,一个就是分析一下AFSecurityPolicy文件,看看AF ...

随机推荐

  1. 四种生成和解析XML文档的方法详解

    众所周知,现在解析XML的方法越来越多,但主流的方法也就四种,即:DOM.SAX.JDOM和DOM4J 下面首先给出这四种方法的jar包下载地址 DOM:在现在的Java JDK里都自带了,在xml- ...

  2. Tomcat的缺省是多少，怎么修改

     Tomcat的缺省端口号是8080. 修改Tomcat端口号: 1.找到Tomcat目录下的conf文件夹 2.进入conf文件夹里面找到server.xml文件 3.打开server.xml文件 ...

  3. AngularJs 隔离作用域

    初学NG,有诸多的不解,今天看了一篇文章,原文地址:https://segmentfault.com/a/1190000002773689#articleHeader0 ,本文运行的代码也出处此. 里 ...

  4. geth常用指令

    ubuntu下载: https://github.com/ethereum/go-ethereum/wiki/Installation-Instructions-for-Ubuntu sudo apt ...

  5. 选择Web框架的20条标准

    原文观点由Matt Raible提出,关于Matt Rabile的介绍:http://www.infoq.com/cn/author/Matt-Raible 内容摘自<Java程序员修炼之道&g ...

  6. python flask中的代码约定

    在Python社区中有许多关于代码风格的约定.如果你写过一段时间Python了,那么也许对此已经有些了解. 我会简单介绍一下,同时给你一些URL链接,从中你可以找到关于这个话题的详细信息. 让我们提出 ...

  7. spring中配置quartz调用两次及项目日志log4j不能每天生成日志解决方法

    在quartz中配置了一个方法运行时会连续调用两次,是因为加载两次,只需在tomcat的server.xml中修改配置 <Host name="www.xx.cn" appB ...

  8. c# 多线程编程中AutoResetEvent和ManualResetEvent

    作为等同于Java的wait,notify,notifyAll的存在,AutoResetEvent和ManualResetEvent分别实现了notify和notifyAll的功能,下面的代码简单讲解 ...

  9. 前端工程化（二）---webpack配置

    导航 前端工程化(一)---工程基础目录搭建 前端工程化(二)---webpack配置 前端工程化(三)---Vue的开发模式 前端工程化(四)---helloWord 继续上一遍的配置,本节主要记录 ...

  10. 前端开发APP，从HBuilder开始~ 【转】

    内容简介 介绍目前前端人员开发app的几种方法,具体介绍hbuilder开发app,一扇赞新的大门~ 无所不能的js 最开始js仅仅局限于网页上一些效果,操作网页内容等, 但是nodejs把js带入了 ...