Golang data.Message函数代码示例

func TestSendConn(t *testing.T) {
	a, b, err := USocketPair()
	if err != nil {
		t.Fatal(err)
	}
	defer a.Close()
	defer b.Close()
	go func() {
		conn, err := SendConn(a, data.Empty().Set("type", "connection").Bytes())
		if err != nil {
			t.Fatal(err)
		}
		if err := conn.Send(data.Empty().Set("foo", "bar").Bytes(), nil); err != nil {
			t.Fatal(err)
		}
		conn.CloseWrite()
	}()
	payload, conn, err := ReceiveConn(b)
	if err != nil {
		t.Fatal(err)
	}
	if val := data.Message(string(payload)).Get("type"); val == nil || val[0] != "connection" {
		t.Fatalf("%v != %v\n", val, "connection")
	}
	msg, _, err := conn.Receive()
	if err != nil {
		t.Fatal(err)
	}
	if val := data.Message(string(msg)).Get("foo"); val == nil || val[0] != "bar" {
		t.Fatalf("%v != %v\n", val, "bar")
	}
}

func (s *Sender) Handle(job *Job) Status {
	cmd := append([]string{job.Name}, job.Args...)
	env := data.Encode(job.Env().MultiMap())
	msg := data.Empty().Set("cmd", cmd...).Set("env", env)
	peer, err := beam.SendConn(s, msg.Bytes())
	if err != nil {
		return job.Errorf("beamsend: %v", err)
	}
	defer peer.Close()
	var tasks sync.WaitGroup
	defer tasks.Wait()
	r := beam.NewRouter(nil)
	r.NewRoute().KeyStartsWith("cmd", "log", "stdout").HasAttachment().Handler(func(p []byte, stdout *os.File) error {
		tasks.Add(1)
		go func() {
			io.Copy(job.Stdout, stdout)
			stdout.Close()
			tasks.Done()
		}()
		return nil
	})
	r.NewRoute().KeyStartsWith("cmd", "log", "stderr").HasAttachment().Handler(func(p []byte, stderr *os.File) error {
		tasks.Add(1)
		go func() {
			io.Copy(job.Stderr, stderr)
			stderr.Close()
			tasks.Done()
		}()
		return nil
	})
	r.NewRoute().KeyStartsWith("cmd", "log", "stdin").HasAttachment().Handler(func(p []byte, stdin *os.File) error {
		go func() {
			io.Copy(stdin, job.Stdin)
			stdin.Close()
		}()
		return nil
	})
	var status int
	r.NewRoute().KeyStartsWith("cmd", "status").Handler(func(p []byte, f *os.File) error {
		cmd := data.Message(p).Get("cmd")
		if len(cmd) != 2 {
			return fmt.Errorf("usage: %s <0-127>", cmd[0])
		}
		s, err := strconv.ParseUint(cmd[1], 10, 8)
		if err != nil {
			return fmt.Errorf("usage: %s <0-127>", cmd[0])
		}
		status = int(s)
		return nil

	})
	if _, err := beam.Copy(r, peer); err != nil {
		return job.Errorf("%v", err)
	}
	return Status(status)
}

func (r *Route) KeyIncludes(k, v string) *Route {
	r.rules = append(r.rules, func(payload []byte, attachment *os.File) bool {
		for _, val := range data.Message(payload).Get(k) {
			if val == v {
				return true
			}
		}
		return false
	})
	return r
}

func CmdLogger(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	if err := os.MkdirAll("logs", 0700); err != nil {
		fmt.Fprintf(stderr, "%v\n", err)
		return
	}
	var tasks sync.WaitGroup
	defer tasks.Wait()
	var n int = 1
	r := beam.NewRouter(out)
	r.NewRoute().HasAttachment().KeyStartsWith("cmd", "log").Handler(func(payload []byte, attachment *os.File) error {
		tasks.Add(1)
		go func(n int) {
			defer tasks.Done()
			defer attachment.Close()
			var streamname string
			if cmd := data.Message(payload).Get("cmd"); len(cmd) == 1 || cmd[1] == "stdout" {
				streamname = "stdout"
			} else {
				streamname = cmd[1]
			}
			if fromcmd := data.Message(payload).Get("fromcmd"); len(fromcmd) != 0 {
				streamname = fmt.Sprintf("%s-%s", strings.Replace(strings.Join(fromcmd, "_"), "/", "_", -1), streamname)
			}
			logfile, err := os.OpenFile(path.Join("logs", fmt.Sprintf("%d-%s", n, streamname)), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0700)
			if err != nil {
				fmt.Fprintf(stderr, "%v\n", err)
				return
			}
			defer logfile.Close()
			io.Copy(logfile, attachment)
			logfile.Sync()
		}(n)
		n++
		return nil
	}).Tee(out)
	if _, err := beam.Copy(r, in); err != nil {
		fmt.Fprintf(stderr, "%v\n", err)
		return
	}
}

func CmdExec(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	cmd := exec.Command(args[1], args[2:]...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	//cmd.Stdin = os.Stdin
	local, remote, err := beam.SocketPair()
	if err != nil {
		fmt.Fprintf(stderr, "%v\n", err)
		return
	}
	child, err := beam.FileConn(local)
	if err != nil {
		local.Close()
		remote.Close()
		fmt.Fprintf(stderr, "%v\n", err)
		return
	}
	local.Close()
	cmd.ExtraFiles = append(cmd.ExtraFiles, remote)

	var tasks sync.WaitGroup
	tasks.Add(1)
	go func() {
		defer Debugf("done copying to child\n")
		defer tasks.Done()
		defer child.CloseWrite()
		beam.Copy(child, in)
	}()

	tasks.Add(1)
	go func() {
		defer Debugf("done copying from child %d\n")
		defer tasks.Done()
		r := beam.NewRouter(out)
		r.NewRoute().All().Handler(func(p []byte, a *os.File) error {
			return out.Send(data.Message(p).Set("pid", fmt.Sprintf("%d", cmd.Process.Pid)).Bytes(), a)
		})
		beam.Copy(r, child)
	}()
	execErr := cmd.Run()
	// We can close both ends of the socket without worrying about data stuck in the buffer,
	// because unix socket writes are fully synchronous.
	child.Close()
	tasks.Wait()
	var status string
	if execErr != nil {
		status = execErr.Error()
	} else {
		status = "ok"
	}
	out.Send(data.Empty().Set("status", status).Set("cmd", args...).Bytes(), nil)
}

func CmdTrace(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	r := beam.NewRouter(out)
	r.NewRoute().All().Handler(func(payload []byte, attachment *os.File) error {
		var sfd string = "nil"
		if attachment != nil {
			sfd = fmt.Sprintf("%d", attachment.Fd())
		}
		fmt.Printf("===> %s [%s]\n", data.Message(payload).Pretty(), sfd)
		out.Send(payload, attachment)
		return nil
	})
	beam.Copy(r, in)
}

func ReceiveFromConn(connections chan net.Conn, dst beam.Sender) error {
	for conn := range connections {
		err := func() error {
			Logf("parsing message from network...\n")
			defer Logf("done parsing message from network\n")
			buf := make([]byte, 4098)
			n, err := conn.Read(buf)
			if n == 0 {
				conn.Close()
				if err == io.EOF {
					return nil
				} else {
					return err
				}
			}
			Logf("decoding message from '%s'\n", buf[:n])
			header, skip, err := data.DecodeString(string(buf[:n]))
			if err != nil {
				conn.Close()
				return err
			}
			pub, priv, err := beam.SocketPair()
			if err != nil {
				return err
			}
			Logf("decoded message: %s\n", data.Message(header).Pretty())
			go func(skipped []byte, conn net.Conn, f *os.File) {
				// this closes both conn and f
				if len(skipped) > 0 {
					if _, err := f.Write(skipped); err != nil {
						Logf("ERROR: %v\n", err)
						f.Close()
						conn.Close()
						return
					}
				}
				bicopy(conn, f)
			}(buf[skip:n], conn, pub)
			if err := dst.Send([]byte(header), priv); err != nil {
				return err
			}
			return nil
		}()
		if err != nil {
			Logf("Error reading from connection: %v\n", err)
		}
	}
	return nil
}

func (r *Route) KeyEquals(k string, full ...string) *Route {
	r.rules = append(r.rules, func(payload []byte, attachment *os.File) bool {
		values := data.Message(payload).Get(k)
		if len(values) != len(full) {
			return false
		}
		for i, v := range full {
			if v != values[i] {
				return false
			}
		}
		return true
	})
	return r
}

func (rcv *Receiver) Run() error {
	r := beam.NewRouter(nil)
	r.NewRoute().KeyExists("cmd").Handler(func(p []byte, f *os.File) error {
		// Use the attachment as a beam return channel
		peer, err := beam.FileConn(f)
		if err != nil {
			f.Close()
			return err
		}
		f.Close()
		defer peer.Close()
		msg := data.Message(p)
		cmd := msg.Get("cmd")
		job := rcv.Engine.Job(cmd[0], cmd[1:]...)
		// Decode env
		env, err := data.Decode(msg.GetOne("env"))
		if err != nil {
			return fmt.Errorf("error decoding 'env': %v", err)
		}
		job.Env().InitMultiMap(env)
		stdout, err := beam.SendRPipe(peer, data.Empty().Set("cmd", "log", "stdout").Bytes())
		if err != nil {
			return err
		}
		job.Stdout.Add(stdout)
		stderr, err := beam.SendRPipe(peer, data.Empty().Set("cmd", "log", "stderr").Bytes())
		if err != nil {
			return err
		}
		job.Stderr.Add(stderr)
		stdin, err := beam.SendWPipe(peer, data.Empty().Set("cmd", "log", "stdin").Bytes())
		if err != nil {
			return err
		}
		job.Stdin.Add(stdin)
		// ignore error because we pass the raw status
		job.Run()
		err = peer.Send(data.Empty().Set("cmd", "status", fmt.Sprintf("%d", job.status)).Bytes(), nil)
		if err != nil {
			return err
		}
		return nil
	})
	_, err := beam.Copy(r, rcv.peer)
	return err
}

func (r *Route) KeyStartsWith(k string, beginning ...string) *Route {
	r.rules = append(r.rules, func(payload []byte, attachment *os.File) bool {
		values := data.Message(payload).Get(k)
		if values == nil {
			return false
		}
		if len(values) < len(beginning) {
			return false
		}
		for i, v := range beginning {
			if v != values[i] {
				return false
			}
		}
		return true
	})
	return r
}

func CmdMultiprint(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	var tasks sync.WaitGroup
	defer tasks.Wait()
	r := beam.NewRouter(out)
	multiprint := func(p []byte, a *os.File) error {
		tasks.Add(1)
		go func() {
			defer tasks.Done()
			defer a.Close()
			msg := data.Message(string(p))
			input := bufio.NewScanner(a)
			for input.Scan() {
				fmt.Printf("[%s] %s\n", msg.Pretty(), input.Text())
			}
		}()
		return nil
	}
	r.NewRoute().KeyIncludes("type", "job").Passthrough(out)
	r.NewRoute().HasAttachment().Handler(multiprint).Tee(out)
	beam.Copy(r, in)
}

func CmdPrint(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	for {
		payload, a, err := in.Receive()
		if err != nil {
			return
		}
		// Skip commands
		if a != nil && data.Message(payload).Get("cmd") == nil {
			dup, err := beam.SendRPipe(out, payload)
			if err != nil {
				a.Close()
				return
			}
			io.Copy(io.MultiWriter(os.Stdout, dup), a)
			dup.Close()
		} else {
			if err := out.Send(payload, a); err != nil {
				return
			}
		}
	}
}

func (rcv *Receiver) Run() error {
	r := beam.NewRouter(nil)
	r.NewRoute().KeyExists("cmd").Handler(func(p []byte, f *os.File) error {
		// Use the attachment as a beam return channel
		peer, err := beam.FileConn(f)
		if err != nil {
			f.Close()
			return err
		}
		cmd := data.Message(p).Get("cmd")
		job := rcv.Engine.Job(cmd[0], cmd[1:]...)
		stdout, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stdout").Bytes())
		if err != nil {
			return err
		}
		job.Stdout.Add(stdout)
		stderr, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stderr").Bytes())
		if err != nil {
			return err
		}
		job.Stderr.Add(stderr)
		stdin, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stdin").Bytes())
		if err != nil {
			return err
		}
		job.Stdin.Add(stdin)
		// ignore error because we pass the raw status
		job.Run()
		err = peer.Send(data.Empty().Set("cmd", "status", fmt.Sprintf("%d", job.status)).Bytes(), nil)
		if err != nil {
			return err
		}
		return nil
	})
	_, err := beam.Copy(r, rcv.peer)
	return err
}

func (r *Route) KeyExists(k string) *Route {
	r.rules = append(r.rules, func(payload []byte, attachment *os.File) bool {
		return data.Message(payload).Get(k) != nil
	})
	return r
}

func (r *Route) NoKey(k string) *Route {
	r.rules = append(r.rules, func(payload []byte, attachment *os.File) bool {
		return len(data.Message(payload).Get(k)) == 0
	})
	return r
}