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. ubuntu下无法编译ruby-2.1.5提示something wrong with CFLAGS -arch x86_64

    在Mac OS X10.10下以下语句运行没有问题: ./configure -prefix=/Users/apple/src/ruby_src/ruby2.1.5_installed --with- ...

  2. .net framework 4 线程安全概述

    线程安全:如果你的代码所在的进程中有多个线程在同时运行,而这些线程可能会同时运行这段代码.如果每次运行结果和单线程运行的结果是一样的,而且其他的变量的值也和预期的是一样的,就是线程安全的.早期的时候, ...

  3. ASP.NET Core 2.0 : 九.从Windows发布到CentOS的跨平台部署

    本文聊一下如何在Windows上用VS开发并发布, 然后将其部署到CentOS上.对于我们一些常在Windows上逛的来说,CentOS用起来还真有些麻烦.MSDN官方有篇文章大概讲了一下(链接),按 ...

  4. c# https请求

    遇到Https网站,c# http请求的时候,总是报SSL连接错误.后来经搜索,发现有解决方案: .net 2.0  需要引入一个第三方组件:BouncyCastle.dll,这是我写的一个例子: p ...

  5. 《转》Xcode 6 正式版如何创建一个Empty Application

    Xcode 6 正式版里面没有Empty Application这个模板,这对于习惯了纯代码编写UI界面的程序员来说很不习惯. 有网友给出了一个解决方法是,把Xcode 6 beta版里面的模板复制过 ...

  6. Spring Cloud入门教程-Hystrix断路器实现容错和降级

    简介 Spring cloud提供了Hystrix容错库用以在服务不可用时,对配置了断路器的方法实行降级策略,临时调用备用方法.这篇文章将创建一个产品微服务,注册到eureka服务注册中心,然后我们使 ...

  7. memocache 分布式搭建

    memcached+magent实现memcached集群   首先说明下memcached存在如下问题 本身没有内置分布式功能,无法实现使用多台Memcache服务器来存储不同的数据,最大程度的使用 ...

  8. 安装 Anaconda 的正确姿势

    下面以 Anaconda2 安装为例, 说明如何更加流畅的使用 Conda Install Anaconda2 安装 Anaconda2(从清华源下载比较快) wget https://mirrors ...

  9. spring 整合 mybatis 中数据源的几种配置方式

    因为spring 整合mybatis的过程中, 有好几种整合方式,尤其是数据源那块,经常看到不一样的配置方式,总感觉有点乱,所以今天有空总结下. 一.采用org.mybatis.spring.mapp ...

  10. saltstack安装部署以及简单实用

    一,saltstack简介:     SaltStack是一种新的基础设施管理方法开发软件,简单易部署,可伸缩的足以管理成千上万的服务器,和足够快的速度控制,与他们交流,以毫秒为单位. SaltSta ...