Golang scope.Context类代码示例

func Serve(ctx scope.Context, addr string) {
	http.Handle("/metrics", prometheus.Handler())

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		ctx.Terminate(err)
	}

	closed := false
	m := sync.Mutex{}
	closeListener := func() {
		m.Lock()
		if !closed {
			listener.Close()
			closed = true
		}
		m.Unlock()
	}

	// Spin off goroutine to watch ctx and close listener if shutdown requested.
	go func() {
		<-ctx.Done()
		closeListener()
	}()

	if err := http.Serve(listener, nil); err != nil {
		fmt.Printf("http[%s]: %s\n", addr, err)
		ctx.Terminate(err)
	}

	closeListener()
	ctx.WaitGroup().Done()
}

func (cmd *serveEmbedCmd) run(ctx scope.Context, args []string) error {
	listener, err := net.Listen("tcp", cmd.addr)
	if err != nil {
		return err
	}

	closed := false
	m := sync.Mutex{}
	closeListener := func() {
		m.Lock()
		if !closed {
			listener.Close()
			closed = true
		}
		m.Unlock()
	}

	// Spin off goroutine to watch ctx and close listener if shutdown requested.
	go func() {
		<-ctx.Done()
		closeListener()
	}()

	if err := http.Serve(listener, cmd); err != nil {
		fmt.Printf("http[%s]: %s\n", cmd.addr, err)
		return err
	}

	closeListener()
	ctx.WaitGroup().Done()
	return ctx.Err()
}

func (cmd *retentionCmd) run(ctx scope.Context, args []string) error {
	heim, b, err := getHeimWithPsqlBackend(ctx)
	if err != nil {
		return err
	}

	defer func() {
		ctx.Cancel()
		ctx.WaitGroup().Wait()
		heim.Backend.Close()
	}()

	// start metrics server
	ctx.WaitGroup().Add(1)
	go retention.Serve(ctx, cmd.addr)

	// start metrics scanner
	ctx.WaitGroup().Add(1)
	go retention.ExpiredScanLoop(ctx, heim.Cluster, b, cmd.interval)

	// start delete scanner
	ctx.WaitGroup().Add(1)
	retention.DeleteScanLoop(ctx, heim.Cluster, b, cmd.interval)

	return nil
}

func (c *Controller) processOne(ctx scope.Context) error {
	job, err := c.claimOrSteal(ctx.ForkWithTimeout(StatsInterval))
	if err != nil {
		if err == scope.TimedOut {
			return nil
		}
		return err
	}

	if job.Type != jobs.EmailJobType {
		return jobs.ErrInvalidJobType
	}

	payload, err := job.Payload()
	if err != nil {
		return err
	}

	return job.Exec(ctx, func(ctx scope.Context) error {
		labels := prometheus.Labels{"queue": c.jq.Name()}
		defer processedCounter.With(labels).Inc()
		if err := c.w.Work(ctx, job, payload); err != nil {
			failedCounter.With(labels).Inc()
			return err
		}
		completedCounter.With(labels).Inc()
		return nil
	})
}

func (jq *JobQueueBinding) WaitForJob(ctx scope.Context) error {
	ch := make(chan error)

	// background goroutine to wait on condition
	go func() {
		// synchronize with caller
		<-ch
		jq.m.Unlock()
		jq.Backend.jobQueueListener().wait(jq.Name())
		ch <- nil
	}()

	// synchronize with background goroutine
	jq.m.Lock()
	ch <- nil
	jq.m.Lock()
	jq.m.Unlock()

	select {
	case <-ctx.Done():
		jq.Backend.jobQueueListener().wakeAll(jq.Name())
		<-ch
		return ctx.Err()
	case err := <-ch:
		return err
	}
}

func (cmd *retentionCmd) run(ctx scope.Context, args []string) error {
	c, err := getCluster(ctx)
	if err != nil {
		return err
	}

	b, err := getBackend(ctx, c)
	if err != nil {
		return fmt.Errorf("backend error: %s", err)
	}
	defer b.Close()

	defer func() {
		ctx.Cancel()
		ctx.WaitGroup().Wait()
	}()

	// start metrics server
	ctx.WaitGroup().Add(1)
	go retention.Serve(ctx, cmd.addr)

	// start metrics scanner
	ctx.WaitGroup().Add(1)
	go retention.ExpiredScanLoop(ctx, c, b, cmd.interval)

	// start delete scanner
	ctx.WaitGroup().Add(1)
	retention.DeleteScanLoop(ctx, c, b, cmd.interval)

	return nil
}

func controller(
	ctx scope.Context, addr string, b proto.Backend, kms security.KMS, c cluster.Cluster) error {

	if addr != "" {
		ctrl, err := console.NewController(ctx, addr, b, kms, c)
		if err != nil {
			return err
		}

		if backend.Config.Console.HostKey != "" {
			if err := ctrl.AddHostKey(backend.Config.Console.HostKey); err != nil {
				return err
			}
		} else {
			if err := ctrl.AddHostKeyFromCluster(backend.Config.Cluster.ServerID); err != nil {
				return err
			}
		}

		for _, authKey := range backend.Config.Console.AuthKeys {
			if authKey == "" {
				continue
			}
			if err := ctrl.AddAuthorizedKeys(authKey); err != nil {
				return err
			}
		}

		ctx.WaitGroup().Add(1)
		go ctrl.Serve()
	}
	return nil
}

// Starts a new bot, blocking until the bot finishes.
func (bcfg *BotConfig) Start(ctx scope.Context, heimDialer proto.ConnectionDialer) error {
	forkedCtx := ctx.Fork()
	conn, err := RetryingConnectionDialer(heimDialer).Dial(forkedCtx, bcfg.RoomName)
	if err != nil {
		return err
	}
	sessions := make(map[string]*Session)

	bot := &Bot{
		ctx:      forkedCtx,
		conn:     conn,
		sessions: sessions,
		nick:     bcfg.Nick,
	}

	for _, cmd := range bcfg.Cmds {
		err := bot.NewSession(RetryingPluginDialer(&cmd))
		if err != nil {
			// If we can't make a session we just failed at start up
			bot.Kill(err)
			return err
		}
	}

	go bot.forwardPacketLoop()

	_ = <-bot.ctx.Done()
	return bot.ctx.Err()
}

func (j *Job) Exec(ctx scope.Context, f func(scope.Context) error) error {
	if j.JobClaim == nil {
		return ErrJobNotClaimed
	}

	w := io.MultiWriter(os.Stdout, j)
	prefix := fmt.Sprintf("[%s-%s] ", j.Queue.Name(), j.HandlerID)
	deadline := time.Now().Add(j.MaxWorkDuration)
	child := logging.LoggingContext(ctx.ForkWithTimeout(j.MaxWorkDuration), w, prefix)
	if err := f(child); err != nil {
		if err != scope.TimedOut {
			delay := time.Duration(j.AttemptsMade+1) * BackoffDuration
			if time.Now().Add(delay).After(deadline) {
				delay = deadline.Sub(time.Now())
			}
			time.Sleep(delay)
		}
		if ferr := j.Fail(ctx, err.Error()); ferr != nil {
			return ferr
		}
		return err
	}

	if err := j.Complete(ctx); err != nil {
		return err
	}

	return nil
}

func (c *Controller) background(ctx scope.Context) {
	defer ctx.WaitGroup().Done()

	var lastStatCheck time.Time
	for {
		logging.Logger(ctx).Printf("[%s] background loop", c.w.QueueName())
		if time.Now().Sub(lastStatCheck) > StatsInterval {
			logging.Logger(ctx).Printf("[%s] collecting stats", c.w.QueueName())
			stats, err := c.jq.Stats(ctx)
			if err != nil {
				logging.Logger(ctx).Printf("error: %s stats: %s", c.w.QueueName(), err)
				return
			}
			lastStatCheck = time.Now()
			labels := map[string]string{"queue": c.w.QueueName()}
			claimedGauge.With(labels).Set(float64(stats.Claimed))
			dueGauge.With(labels).Set(float64(stats.Due))
			waitingGauge.With(labels).Set(float64(stats.Waiting))
		}
		if err := c.processOne(ctx); err != nil {
			// TODO: retry a couple times before giving up
			logging.Logger(ctx).Printf("error: %s: %s", c.w.QueueName(), err)
			return
		}
	}
}

func (et *EmailTracker) Send(
	ctx scope.Context, js jobs.JobService, templater *templates.Templater, deliverer emails.Deliverer,
	account proto.Account, to, templateName string, data interface{}) (
	*emails.EmailRef, error) {

	if to == "" {
		to, _ = account.Email()
	}

	sf, err := snowflake.New()
	if err != nil {
		return nil, err
	}
	msgID := fmt.Sprintf("<%[email protected]%s>", sf, deliverer.LocalName())

	ref, err := emails.NewEmail(templater, msgID, to, templateName, data)
	if err != nil {
		return nil, err
	}
	ref.AccountID = account.ID()

	jq, err := js.GetQueue(ctx, jobs.EmailQueue)
	if err != nil {
		return nil, err
	}

	payload := &jobs.EmailJob{
		AccountID: account.ID(),
		EmailID:   ref.ID,
	}
	job, err := jq.AddAndClaim(ctx, jobs.EmailJobType, payload, "immediate", jobs.EmailJobOptions...)
	if err != nil {
		return nil, err
	}

	ref.JobID = job.ID

	et.m.Lock()
	if et.emailsByAccount == nil {
		et.emailsByAccount = map[snowflake.Snowflake][]*emails.EmailRef{}
	}
	et.emailsByAccount[account.ID()] = append(et.emailsByAccount[account.ID()], ref)
	et.m.Unlock()

	child := ctx.Fork()
	child.WaitGroup().Add(1)

	go job.Exec(child, func(ctx scope.Context) error {
		defer ctx.WaitGroup().Done()

		logging.Logger(ctx).Printf("delivering to %s\n", to)
		if err := deliverer.Deliver(ctx, ref); err != nil {
			return err
		}
		return nil
	})

	return ref, nil
}

func (c *Client) FromRequest(ctx scope.Context, r *http.Request) {
	c.UserAgent = r.Header.Get("User-Agent")
	c.Connected = time.Now()
	c.IP = getIP(r)

	var k clientKey
	ctx.Set(k, c)
}

func (c *Client) FromContext(ctx scope.Context) bool {
	var k clientKey
	src, ok := ctx.Get(k).(*Client)
	if !ok || src == nil {
		return false
	}
	*c = *src
	return true
}

func (d *RetryPluginDialer) Dial(ctx scope.Context) (proto.Plugin, error) {
	retry := &RetryPlugin{
		ctx:    ctx.Fork(),
		dialer: d.Dialer,
	}
	err := retry.start(nil)
	if err != nil {
		return nil, err
	}
	return retry, err
}

func (s *session) Send(ctx scope.Context, cmdType proto.PacketType, payload interface{}) error {
	// Special case: certain events have privileged info that may need to be stripped from them
	switch event := payload.(type) {
	case *proto.PresenceEvent:
		switch s.privilegeLevel() {
		case proto.Staff:
		case proto.Host:
			event.RealClientAddress = ""
		default:
			event.RealClientAddress = ""
			event.ClientAddress = ""
		}
	case *proto.Message:
		if s.privilegeLevel() == proto.General {
			event.Sender.ClientAddress = ""
		}
	case *proto.EditMessageEvent:
		if s.privilegeLevel() == proto.General {
			event.Sender.ClientAddress = ""
		}
	}

	var err error
	payload, err = proto.DecryptPayload(payload, &s.client.Authorization, s.privilegeLevel())
	if err != nil {
		return err
	}

	encoded, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	cmd := &proto.Packet{
		Type: cmdType,
		Data: encoded,
	}

	// Add to outgoing channel. If channel is full, defer to goroutine so as not to block
	// the caller (this may result in deliveries coming out of order).
	select {
	case <-ctx.Done():
		// Session is closed, return error.
		return ctx.Err()
	case s.outgoing <- cmd:
		// Packet delivered to queue.
	default:
		// Queue is full.
		logging.Logger(s.ctx).Printf("outgoing channel full, ordering cannot be guaranteed")
		go func() { s.outgoing <- cmd }()
	}

	return nil
}