memlimiter/service_impl.go at master · newcloudtechnologies/memlimiter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
/*
 * Copyright (c) New Cloud Technologies, Ltd. 2013-2026.
 * Author: Vitaly Isaev <vitaly.isaev@myoffice.team>
 * License: https://github.com/newcloudtechnologies/memlimiter/blob/master/LICENSE
 */

package memlimiter

import (
	"errors"
	"fmt"
	"math"
	"runtime/debug"

	"github.com/go-logr/logr"
	"github.com/newcloudtechnologies/memlimiter/backpressure"
	"github.com/newcloudtechnologies/memlimiter/controller"
	"github.com/newcloudtechnologies/memlimiter/controller/nextgc"
	"github.com/newcloudtechnologies/memlimiter/middleware"
	"github.com/newcloudtechnologies/memlimiter/stats"
	"github.com/newcloudtechnologies/memlimiter/utils/config/prepare"
)

var _ Service = (*serviceImpl)(nil)

type serviceImpl struct {
	middleware           middleware.Middleware
	backpressureOperator backpressure.Operator
	statsSubscription    stats.ServiceStatsSubscription
	controller           controller.Controller
	restoreGoMemoryLimit bool
	oldGoMemoryLimit     int64
	logger               logr.Logger
}

func (s *serviceImpl) Middleware() middleware.Middleware { return s.middleware }

func (s *serviceImpl) GetStats() (*stats.MemLimiterStats, error) {
	controllerStats, err := s.controller.GetStats()
	if err != nil {
		return nil, fmt.Errorf("controller tracker: %w", err)
	}

	backpressureStats, err := s.backpressureOperator.GetStats()
	if err != nil {
		return nil, fmt.Errorf("backpressure tracker: %w", err)
	}

	return &stats.MemLimiterStats{
		Controller:   controllerStats,
		Backpressure: backpressureStats,
	}, nil
}

func (s *serviceImpl) Quit() {
	s.logger.Info("terminating MemLimiter service")
	s.controller.Quit()
	s.statsSubscription.Quit()
	s.backpressureOperator.Quit()

	if s.restoreGoMemoryLimit {
		debug.SetMemoryLimit(s.oldGoMemoryLimit)
	}
}

// newServiceImpl - main entrypoint for MemLimiter.
func newServiceImpl(
	logger logr.Logger,
	cfg *Config,
	statsSubscription stats.ServiceStatsSubscription,
	backpressureOperator backpressure.Operator,
) (Service, error) {
	if err := prepare.Prepare(cfg); err != nil {
		return nil, fmt.Errorf("prepare config: %w", err)
	}

	if statsSubscription == nil {
		return nil, errors.New("nil tracker subscription passed")
	}

	var (
		restoreGoMemoryLimit bool
		oldGoMemoryLimit     int64
	)

	if cfg.GoMemoryLimit.Value > 0 {
		if cfg.GoMemoryLimit.Value > uint64(math.MaxInt64) {
			return nil, errors.New("go memory limit exceeds int64 range")
		}

		oldGoMemoryLimit = debug.SetMemoryLimit(int64(cfg.GoMemoryLimit.Value))
		restoreGoMemoryLimit = true
	}

	logger.Info("starting MemLimiter service")

	c, err := nextgc.NewControllerFromConfig(
		logger,
		cfg.ControllerNextGC,
		statsSubscription,
		backpressureOperator,
	)
	if err != nil {
		if restoreGoMemoryLimit {
			debug.SetMemoryLimit(oldGoMemoryLimit)
		}

		return nil, fmt.Errorf("new controller from config: %w", err)
	}

	return &serviceImpl{
		middleware:           middleware.NewMiddleware(logger, backpressureOperator),
		backpressureOperator: backpressureOperator,
		statsSubscription:    statsSubscription,
		controller:           c,
		restoreGoMemoryLimit: restoreGoMemoryLimit,
		oldGoMemoryLimit:     oldGoMemoryLimit,
		logger:               logger,
	}, nil
}