apkipa 1 week ago
parent
commit
a37395028c
11 changed files with 632 additions and 115 deletions
  1. 3 0
      .gitmodules
  2. 3 0
      config/application.yaml
  3. 9 1
      go.mod
  4. 11 2
      go.sum
  5. 388 111
      main.go
  6. 1 0
      minio-into-stck-installer/update.sh
  7. 2 1
      package.sh
  8. 1 0
      stck-nsq-msg
  9. 94 0
      util/dchan.go
  10. 72 0
      util/file_util.go
  11. 48 0
      util/util.go

+ 3 - 0
.gitmodules

@@ -0,0 +1,3 @@
+[submodule "stck-nsq-msg"]
+	path = stck-nsq-msg
+	url = http://47.101.131.235:3000/chenyun3/stck-nsq-msg

+ 3 - 0
config/application.yaml

@@ -12,6 +12,9 @@ stck:
   host: "172.31.48.206:9000"
   # host: "localhost:9001"
   table: "tsdb_cpp_dist"
+nsq:
+  tcpAddr: "localhost:4150"
+  lookupdHttpAddr: "localhost:4161"
 main:
   uploadRetryMaxTimes: 20
   failedRetryDelaySeconds: 5

+ 9 - 1
go.mod

@@ -2,28 +2,36 @@ module example/minio-into-stck
 
 go 1.22.2
 
+replace stck/stck-nsq-msg => ./stck-nsq-msg
+
 require (
 	github.com/ClickHouse/clickhouse-go/v2 v2.30.0
 	github.com/fsnotify/fsnotify v1.7.0
 	github.com/klauspost/compress v1.17.11
 	github.com/minio/minio-go/v7 v7.0.80
 	github.com/natefinch/lumberjack v2.0.0+incompatible
+	github.com/nsqio/go-nsq v1.1.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/viper v1.19.0
 	gorm.io/driver/mysql v1.5.7
 	gorm.io/gorm v1.25.12
+	stck/stck-nsq-msg v0.0.0-00010101000000-000000000000
 )
 
 require (
 	github.com/BurntSushi/toml v1.4.0 // indirect
 	github.com/ClickHouse/ch-go v0.61.5 // indirect
+	github.com/Masterminds/semver/v3 v3.3.1 // indirect
 	github.com/andybalholm/brotli v1.1.1 // indirect
+	github.com/denisbrodbeck/machineid v1.0.1 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/go-faster/city v1.0.1 // indirect
 	github.com/go-faster/errors v0.7.1 // indirect
 	github.com/go-ini/ini v1.67.0 // indirect
 	github.com/go-sql-driver/mysql v1.7.0 // indirect
+	github.com/go-viper/mapstructure/v2 v2.2.1 // indirect
 	github.com/goccy/go-json v0.10.3 // indirect
+	github.com/golang/snappy v0.0.1 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
@@ -52,7 +60,7 @@ require (
 	golang.org/x/crypto v0.28.0 // indirect
 	golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
 	golang.org/x/net v0.30.0 // indirect
-	golang.org/x/sys v0.26.0 // indirect
+	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/text v0.19.0 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect

+ 11 - 2
go.sum

@@ -4,12 +4,16 @@ github.com/ClickHouse/ch-go v0.61.5 h1:zwR8QbYI0tsMiEcze/uIMK+Tz1D3XZXLdNrlaOpeE
 github.com/ClickHouse/ch-go v0.61.5/go.mod h1:s1LJW/F/LcFs5HJnuogFMta50kKDO0lf9zzfrbl0RQg=
 github.com/ClickHouse/clickhouse-go/v2 v2.30.0 h1:AG4D/hW39qa58+JHQIFOSnxyL46H6h2lrmGGk17dhFo=
 github.com/ClickHouse/clickhouse-go/v2 v2.30.0/go.mod h1:i9ZQAojcayW3RsdCb3YR+n+wC2h65eJsZCscZ1Z1wyo=
+github.com/Masterminds/semver/v3 v3.3.1 h1:QtNSWtVZ3nBfk8mAOu/B6v7FMJ+NHTIgUPi7rj+4nv4=
+github.com/Masterminds/semver/v3 v3.3.1/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
 github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
 github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/denisbrodbeck/machineid v1.0.1 h1:geKr9qtkB876mXguW2X6TU4ZynleN6ezuMSRhl4D7AQ=
+github.com/denisbrodbeck/machineid v1.0.1/go.mod h1:dJUwb7PTidGDeYyUBmXZ2GphQBbjJCrnectwCyxcUSI=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
@@ -24,10 +28,13 @@ github.com/go-ini/ini v1.67.0 h1:z6ZrTEZqSWOTyH2FlglNbNgARyHG8oLW9gMELqKr06A=
 github.com/go-ini/ini v1.67.0/go.mod h1:ByCAeIL28uOIIG0E3PJtZPDL8WnHpFKFOtgjp+3Ies8=
 github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc=
 github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
+github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
+github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/goccy/go-json v0.10.3 h1:KZ5WoDbxAIgm2HNbYckL0se1fHD6rz5j4ywS6ebzDqA=
 github.com/goccy/go-json v0.10.3/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
@@ -67,6 +74,8 @@ github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RR
 github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc=
 github.com/natefinch/lumberjack v2.0.0+incompatible h1:4QJd3OLAMgj7ph+yZTuX13Ld4UpgHp07nNdFX7mqFfM=
 github.com/natefinch/lumberjack v2.0.0+incompatible/go.mod h1:Wi9p2TTF5DG5oU+6YfsmYQpsTIOm0B1VNzQg9Mw6nPk=
+github.com/nsqio/go-nsq v1.1.0 h1:PQg+xxiUjA7V+TLdXw7nVrJ5Jbl3sN86EhGCQj4+FYE=
+github.com/nsqio/go-nsq v1.1.0/go.mod h1:vKq36oyeVXgsS5Q8YEO7WghqidAVXQlcFxzQbQTuDEY=
 github.com/paulmach/orb v0.11.1 h1:3koVegMC4X/WeiXYz9iswopaTwMem53NzTJuTF20JzU=
 github.com/paulmach/orb v0.11.1/go.mod h1:5mULz1xQfs3bmQm63QEJA6lNGujuRafwA5S/EnuLaLU=
 github.com/paulmach/protoscan v0.2.1/go.mod h1:SpcSwydNLrxUGSDvXvO0P7g7AuhJ7lcKfDlhJCDw2gY=
@@ -161,8 +170,8 @@ golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
-golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
+golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

+ 388 - 111
main.go

@@ -8,7 +8,13 @@ import (
 	"example/minio-into-stck/util"
 	"fmt"
 	"os"
+	"os/exec"
+	"os/signal"
+	"path/filepath"
+	"stck/stck-nsq-msg"
+	smsg "stck/stck-nsq-msg/msg"
 	"sync"
+	"syscall"
 
 	"io"
 	// "log"
@@ -26,6 +32,7 @@ import (
 	"github.com/minio/minio-go/v7/pkg/credentials"
 	"github.com/minio/minio-go/v7/pkg/notification"
 	"github.com/natefinch/lumberjack"
+	"github.com/nsqio/go-nsq"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
 	"gorm.io/driver/mysql"
@@ -51,6 +58,10 @@ type AppInitConfig struct {
 		Host  string `mapstructure:"host"`
 		Table string `mapstructure:"table"`
 	} `mapstructure:"stck"`
+	Nsq struct {
+		TcpAddr         string `mapstructure:"tcpAddr"`
+		LookupdHttpAddr string `mapstructure:"lookupdHttpAddr"`
+	} `mapstructure:"nsq"`
 	Main struct {
 		UploadRetryMaxTimes     int `mapstructure:"uploadRetryMaxTimes"`
 		FailedRetryDelaySeconds int `mapstructure:"failedRetryDelaySeconds"`
@@ -87,9 +98,234 @@ type StreamMetadata struct {
 	Name            string `json:"name"`
 	TimestampOffset int64  `json:"timestamp_offset"`
 	Interval        int64  `json:"interval"`
-	PartsCount      int    `json:"parts_count"`
-	PointsPerPart   int    `json:"points_per_part"`
-	TotalPoints     int    `json:"total_points"`
+	PartsCount      int64  `json:"parts_count"`
+	PointsPerPart   int64  `json:"points_per_part"`
+	TotalPoints     int64  `json:"total_points"`
+}
+
+var gLocalIpAddress string
+var gMachineID string
+var gNsqProducer *nsq.Producer
+var gNsqConsumer *nsq.Consumer
+var gAppStartTime time.Time = time.Now()
+var programmaticQuitChan chan struct{} = make(chan struct{}, 1)
+var gAppQuitting = false
+var gAppExitWaitGroup sync.WaitGroup
+
+var buildtime string
+
+func MakeHeartbeatMsg() *smsg.DeviceHeartbeatMsg {
+	m := smsg.MakeDeviceHeartbeatMsg(gMachineID, gAppStartTime.UnixNano(), time.Now().UnixNano(),
+		gLocalIpAddress, "minio-into-stck", "0.1.0+dev."+buildtime)
+	return &m
+}
+
+func PublishMessage(msg stcknsqmsg.StckNsqMsgVariant) error {
+	payload, err := stcknsqmsg.ToStckNsqMsgString(msg)
+	if err != nil {
+		return fmt.Errorf("marshal message error: %w", err)
+	}
+	err = gNsqProducer.Publish(smsg.ServerDoNotifyTopic, []byte(payload))
+	if err != nil {
+		return fmt.Errorf("publish message error: %w", err)
+	}
+	return nil
+}
+
+func initNsq() {
+	ip, err := stcknsqmsg.GetLocalIP()
+	if err != nil {
+		logger.Warnf("GetLocalIP error: %s", err)
+	} else {
+		gLocalIpAddress = ip.String()
+		logger.Infof("Local IP: %s", gLocalIpAddress)
+	}
+	gMachineID = stcknsqmsg.MakeUniqueMachineID()
+	logger.Infof("Machine ID: %s", gMachineID)
+
+	fmt.Printf("IP: %s, Machine ID: %s\n", gLocalIpAddress, gMachineID)
+
+	// Connect to NSQ
+	nsqConfig := nsq.NewConfig()
+	gNsqProducer, err = nsq.NewProducer(appInitCfg.Nsq.TcpAddr, nsqConfig)
+	if err != nil {
+		logger.Fatalf("NSQ Producer init error: %s", err)
+	}
+	gNsqConsumer, err = nsq.NewConsumer(smsg.ClientDoActionTopic, gMachineID, nsqConfig)
+	if err != nil {
+		logger.Fatalf("NSQ Consumer init error: %s", err)
+	}
+	gNsqConsumer.AddConcurrentHandlers(nsq.HandlerFunc(func(message *nsq.Message) error {
+		logger.Debugf("NSQ Consumer received message: %s", message.Body)
+
+		// Parse the message
+		recvTime := message.Timestamp
+		msg, err := stcknsqmsg.FromString(string(message.Body))
+		if err != nil {
+			logger.Errorf("NSQ Consumer unmarshal message error: %s", err)
+			return nil
+		}
+
+		// Process the message
+		switch data := msg.Data.(type) {
+		case *smsg.DevicePingMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			// Write pong
+			pongMsg := smsg.MakeDevicePongMsg(gMachineID, recvTime)
+			err := PublishMessage(&pongMsg)
+			if err != nil {
+				logger.Errorf("send pong error: %s", err)
+			}
+		case *smsg.RequestDeviceExecuteShellScriptMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			// Execute the shell script
+			resp := smsg.MakeDeviceActionDoneMsg(gMachineID, data.MsgType(), "", recvTime, -1, "")
+			result, err := func(data *smsg.RequestDeviceExecuteShellScriptMsg) (string, error) {
+				cmd := exec.Command("bash", "-c", data.Script)
+				var out bytes.Buffer
+				cmd.Stdout = &out
+				cmd.Stderr = &out
+				err := cmd.Run()
+				return out.String(), err
+			}(data)
+			if err != nil {
+				errMsg := fmt.Sprintf("execute shell script error:\n%s\n\noutput:\n%s", err, result)
+				// Write error message
+				resp.Status = -1
+				resp.Msg = errMsg
+			} else {
+				// Write output
+				resp.Status = 0
+				resp.Msg = result
+			}
+			err = PublishMessage(&resp)
+			if err != nil {
+				logger.Errorf("send action done error: %s", err)
+			}
+		case *smsg.RequestDeviceUpdateMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			if data.ServiceName != "minio-into-stck" {
+				break
+			}
+			resp := smsg.MakeDeviceActionDoneMsg(gMachineID, data.MsgType(), "", recvTime, 0, "")
+			result, err := func(data *smsg.RequestDeviceUpdateMsg) (string, error) {
+				// Download the update file to /tmp
+				downloadPath := "/tmp/minio-into-stck-updater"
+				updateScriptPath := filepath.Join(downloadPath, "update.sh")
+				var out bytes.Buffer
+				err := os.RemoveAll(downloadPath)
+				if err != nil {
+					return "", err
+				}
+				err = os.MkdirAll(downloadPath, 0777)
+				if err != nil {
+					return "", err
+				}
+				updateScriptContent := fmt.Sprintf(`#!/bin/bash
+set -e
+cd %s
+wget --tries=3 -nv -O installer.tar.gz %s
+tar -xzf installer.tar.gz
+cd minio-into-stck-installer
+./replacing-update.sh
+`, downloadPath, data.ServiceBinaryURL)
+				err = os.WriteFile(updateScriptPath, []byte(updateScriptContent), 0777)
+				if err != nil {
+					return "", err
+				}
+				// Execute the update script
+				cmd := exec.Command("bash", "-c", updateScriptPath)
+				cmd.Stdout = &out
+				cmd.Stderr = &out
+				err = cmd.Run()
+				return out.String(), err
+			}(data)
+			if err != nil {
+				errMsg := fmt.Sprintf("execute update process error:\n%s\n\noutput:\n%s", err, result)
+				// Write error message
+				resp.Status = -1
+				resp.Msg = errMsg
+			} else {
+				// Write output
+				resp.Status = 0
+				resp.Msg = "executed update process successfully\n\noutput:\n" + result
+			}
+			err = PublishMessage(&resp)
+			if err != nil {
+				logger.Errorf("send action done error: %s", err)
+			}
+		case *smsg.ForceDeviceRebootMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			programmaticQuitChan <- struct{}{}
+		case *smsg.ForceDeviceSendHeartbeatMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			// Send heartbeat
+			heartBeatMsg := MakeHeartbeatMsg()
+			err := PublishMessage(heartBeatMsg)
+			if err != nil {
+				logger.Errorf("send heartbeat error: %s", err)
+			}
+		case *smsg.RequestDeviceUploadLogsToMinioMsg:
+			if !stcknsqmsg.DeviceIdMatches(data.DeviceID, gMachineID) {
+				break
+			}
+			// Upload logs to MinIO
+			resp := smsg.MakeDeviceActionDoneMsg(gMachineID, data.MsgType(), "", recvTime, 0, "")
+			minioCreds := credentials.NewStaticV4(appInitCfg.Minio.AccessKey, appInitCfg.Minio.Secret, "")
+			minioClient, err := minio.New(appInitCfg.Minio.Host, &minio.Options{
+				Creds:  minioCreds,
+				Secure: false,
+			})
+			if err != nil {
+				logger.Errorf("MinIO Client init error: %s", err)
+				resp.Status = -1
+				resp.Msg += fmt.Sprintf("MinIO Client init error: %s\n", err)
+			} else {
+				if util.FileExists("./logs") {
+					err = util.MinioUploadFolder(minioClient, data.RemoteBucket,
+						filepath.Join(data.RemoteBasePath, gMachineID, "logs"), "logs")
+					if err != nil {
+						logger.Errorf("upload logs to MinIO error: %s", err)
+						resp.Status = -1
+						resp.Msg += fmt.Sprintf("upload logs to MinIO error: %s\n", err)
+					}
+				}
+				if util.FileExists("./log") {
+					err = util.MinioUploadFolder(minioClient, data.RemoteBucket,
+						filepath.Join(data.RemoteBasePath, gMachineID, "log"), "log")
+					if err != nil {
+						logger.Errorf("upload log to MinIO error: %s", err)
+						resp.Status = -1
+						resp.Msg += fmt.Sprintf("upload log to MinIO error: %s\n", err)
+					}
+				}
+			}
+			err = PublishMessage(&resp)
+			if err != nil {
+				logger.Errorf("send action done error: %s", err)
+			}
+		default:
+			logger.Debugf("NSQ Consumer ignored unknown or uninteresting message: %v", msg)
+		}
+
+		// Notify NSQ that the message is processed successfully
+		return nil
+	}), 1)
+	// err = gNsqConsumer.ConnectToNSQLookupd(gAppConfig.Nsq.LookupdHttpAddr)
+	err = gNsqConsumer.ConnectToNSQD(appInitCfg.Nsq.TcpAddr)
+	if err != nil {
+		logger.Fatalf("NSQ Consumer connect error: %s", err)
+	}
 }
 
 var appInitCfg *AppInitConfig = &AppInitConfig{}
@@ -202,43 +438,45 @@ var statLogger *logrus.Logger
 var mutexObjFailCounter = &sync.Mutex{}
 var objFailCounter map[string]int = make(map[string]int)
 
-var mutexHasObjEmitted = &sync.Mutex{}
-var hasObjEmitted map[string]bool = make(map[string]bool)
-
-func markObjEmitted(obj string) {
-	mutexHasObjEmitted.Lock()
-	hasObjEmitted[obj] = true
-	mutexHasObjEmitted.Unlock()
-}
-
-func unmarkObjEmitted(obj string) {
-	mutexHasObjEmitted.Lock()
-	delete(hasObjEmitted, obj)
-	mutexHasObjEmitted.Unlock()
-}
-
-func hasObjAlreadyEmitted(obj string) bool {
-	mutexHasObjEmitted.Lock()
-	defer mutexHasObjEmitted.Unlock()
-	value, exists := hasObjEmitted[obj]
-	return exists && value
-}
-
-func writeToObjChanDeduplicated(obj string, objChan chan<- string) {
-	if !hasObjAlreadyEmitted(obj) {
-		markObjEmitted(obj)
-		objChan <- obj
-	}
-}
-
-func readFromObjChanDeduplicated(objChan <-chan string) string {
-	obj := <-objChan
-	unmarkObjEmitted(obj)
-	return obj
-}
+// var mutexHasObjEmitted = &sync.Mutex{}
+// var hasObjEmitted map[string]bool = make(map[string]bool)
+
+// func markObjEmitted(obj string) {
+// 	mutexHasObjEmitted.Lock()
+// 	hasObjEmitted[obj] = true
+// 	mutexHasObjEmitted.Unlock()
+// }
+
+// func unmarkObjEmitted(obj string) {
+// 	mutexHasObjEmitted.Lock()
+// 	delete(hasObjEmitted, obj)
+// 	mutexHasObjEmitted.Unlock()
+// }
+
+// func hasObjAlreadyEmitted(obj string) bool {
+// 	mutexHasObjEmitted.Lock()
+// 	defer mutexHasObjEmitted.Unlock()
+// 	value, exists := hasObjEmitted[obj]
+// 	return exists && value
+// }
+
+// func writeToObjChanDeduplicated(obj string, objChan chan<- string) {
+// 	if !hasObjAlreadyEmitted(obj) {
+// 		markObjEmitted(obj)
+// 		objChan <- obj
+// 	}
+// }
+
+// func readFromObjChanDeduplicated(objChan <-chan string) string {
+// 	obj := <-objChan
+// 	unmarkObjEmitted(obj)
+// 	return obj
+// }
 
 func main() {
-	fmt.Println("Starting application...")
+	var err error
+
+	fmt.Println("Starting minio-into-stck, build time:", buildtime)
 
 	logger = initLog()
 
@@ -247,7 +485,9 @@ func main() {
 	// Load configuration from file
 	initLoadConfig()
 
-	var err error
+	// 初始化 NSQ
+	initNsq()
+
 	var db *gorm.DB
 	var minioClient *minio.Client
 	var ckConn driver.Conn
@@ -317,7 +557,33 @@ func main() {
 
 	// Start the main work
 	logger.Infoln("Starting main worker...")
-	main_worker(AppCtx{db, minioClient, ckConn})
+	gAppExitWaitGroup.Add(1)
+	go func() {
+		defer gAppExitWaitGroup.Done()
+		main_worker(AppCtx{db, minioClient, ckConn})
+	}()
+
+	// Wait on signal.
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+
+	select {
+	case <-signalChan:
+		logger.Infof("received signal, stopping watch-daemon")
+	case <-programmaticQuitChan:
+		logger.Infof("received programmatic quit signal, stopping watch-daemon")
+	}
+
+	gAppQuitting = true
+
+	// Close the NSQ producer and consumer
+	gNsqProducer.Stop()
+	gNsqConsumer.Stop()
+
+	// Wait for the goroutines to exit
+	gAppExitWaitGroup.Wait()
+
+	logger.Infof("watch-daemon stopped gracefully")
 }
 
 type AppCtx struct {
@@ -337,7 +603,8 @@ type PartUploadArgs struct {
 func main_worker(app AppCtx) {
 	ctx := context.Background()
 
-	objUploadChan := make(chan string, 1024*256)
+	// objUploadChan := make(chan string, 1024*256)
+	objUploadChan := util.NewDChan[string](1024 * 16)
 
 	// Load config from DB
 	appCfg, err := load_app_cfg_from_db(app.db)
@@ -355,12 +622,9 @@ func main_worker(app AppCtx) {
 				ctx, appInitCfg.Minio.Bucket, "", "", []string{string(notification.ObjectCreatedAll)})
 
 			// Listen OK, start the full upload trigger to upload maybe missed files
-			go trigger_full_upload(app, objUploadChan)
-
-			lastSentKey := ""
+			go trigger_full_upload(app, objUploadChan.In())
 
 			for notifyInfo := range notifys {
-				lastSentKey = ""
 				for _, record := range notifyInfo.Records {
 					key := record.S3.Object.Key
 					logger.Traceln("New object notification:", key)
@@ -371,11 +635,8 @@ func main_worker(app AppCtx) {
 						continue
 					}
 					key = strings.Join(keyParts[:len(keyParts)-1], "/") + "/"
-					if key != lastSentKey {
-						lastSentKey = key
-						// objUploadChan <- key
-						writeToObjChanDeduplicated(key, objUploadChan)
-					}
+					// Queue the object for upload
+					objUploadChan.Write(key)
 				}
 				if notifyInfo.Err != nil {
 					logger.Errorf("Bucket notification listener error: %v", notifyInfo.Err)
@@ -387,8 +648,13 @@ func main_worker(app AppCtx) {
 	}()
 
 	// Start the main loop (streams upload worker)
-	for objToUpload := range objUploadChan {
-		unmarkObjEmitted(objToUpload)
+	for objToUpload := range objUploadChan.Out() {
+		objUploadChan.MarkElemReadDone(objToUpload)
+
+		if gAppQuitting {
+			logger.Infof("Quitting, stopping main worker")
+			return
+		}
 
 		logger.Infoln("Checking stream object:", objToUpload)
 		if object_is_blacklisted(objToUpload) {
@@ -408,16 +674,12 @@ func main_worker(app AppCtx) {
 		fullyUploaded, err := upload_one_stream(app, objToUpload)
 		if err != nil {
 			// Queue the object for retry
-			logger.Warnf("Failed to upload stream `%s`: `%v`, retrying", objToUpload, err)
+			logger.Warnf("Failed to upload stream `%s`: `%v`, retrying after %d seconds",
+				objToUpload, err, appInitCfg.Main.FailedRetryDelaySeconds)
 			mutexObjFailCounter.Lock()
 			objFailCounter[objToUpload]++
 			mutexObjFailCounter.Unlock()
-			go func() {
-				markObjEmitted(objToUpload)
-				time.Sleep(time.Duration(appInitCfg.Main.FailedRetryDelaySeconds) * time.Second)
-				objUploadChan <- objToUpload
-				// writeToObjChanDeduplicated(objToUpload, objUploadChan)
-			}()
+			objUploadChan.DelayedWrite(objToUpload, time.Duration(appInitCfg.Main.FailedRetryDelaySeconds)*time.Second)
 			continue
 		}
 
@@ -447,12 +709,7 @@ func main_worker(app AppCtx) {
 			mutexObjFailCounter.Unlock()
 			logger.Warnf("Stream %s is not fully uploaded, retrying after %d seconds",
 				objToUpload, appInitCfg.Main.FailedRetryDelaySeconds)
-			go func() {
-				markObjEmitted(objToUpload)
-				time.Sleep(time.Duration(appInitCfg.Main.FailedRetryDelaySeconds) * time.Second)
-				objUploadChan <- objToUpload
-				// writeToObjChanDeduplicated(objToUpload, objUploadChan)
-			}()
+			objUploadChan.DelayedWrite(objToUpload, time.Duration(appInitCfg.Main.FailedRetryDelaySeconds)*time.Second)
 		}
 	}
 }
@@ -486,20 +743,10 @@ func trigger_full_upload(app AppCtx, objToUploadChan chan<- string) {
 					// Is a directory, should be a stream then
 					uploaded := object_already_uploaded(app, key)
 					if !uploaded {
-						markObjEmitted(key)
 						objToUploadChan <- key
-						// writeToObjChanDeduplicated(key, objToUploadChan)
 					}
 				}
 			}
-
-			// folder := strings.Split(key, "/")[0]
-			// uploaded := object_already_uploaded(app, folder)
-			// if !uploaded {
-			// 	markObjEmitted(folder)
-			// 	objToUploadChan <- folder
-			// 	// writeToObjChanDeduplicated(folder, objToUploadChan)
-			// }
 		}
 	}
 }
@@ -513,12 +760,13 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 		Msg       string // "error message"
 	}
 	type StreamUploadStatistics struct {
-		StartTime      time.Time
-		EndTime        time.Time
-		StreamName     string
-		MetaPartsCount int
-		Objects        map[string]StreamObjectUploadStatistics
-		Msg            string
+		StartTime       time.Time
+		EndTime         time.Time
+		StreamName      string
+		MetaPointsCount int64
+		MetaPartsCount  int64
+		Objects         map[string]StreamObjectUploadStatistics
+		Msg             string
 	}
 	streamStats := StreamUploadStatistics{
 		StartTime:  time.Now(),
@@ -535,10 +783,10 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 
 	defer func() {
 		streamStats.EndTime = time.Now()
-		repeatedCount := 0
-		okCount := 0
-		failCount := 0
-		totalCount := 0
+		repeatedCount := int64(0)
+		okCount := int64(0)
+		failCount := int64(0)
+		totalCount := int64(0)
 		objDetailsMsg := ""
 		for _, obj := range streamStats.Objects {
 			totalCount++
@@ -562,6 +810,15 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 			util.FmtMyTime(streamStats.StartTime), util.FmtMyTime(streamStats.EndTime),
 			fullyUploaded, streamStats.MetaPartsCount,
 			repeatedCount, okCount, failCount, totalCount, objDetailsMsg)
+
+		// Send upload status changed message
+		msg := smsg.MakeStreamInsertToStckStatusChangedMsg(gMachineID, streamStats.StartTime.UnixNano(),
+			streamStats.StreamName, streamStats.MetaPointsCount, streamStats.MetaPartsCount,
+			okCount, failCount, repeatedCount)
+		err := PublishMessage(&msg)
+		if err != nil {
+			logger.Errorf("send stream insert to stck status changed message error: %s", err)
+		}
 	}()
 
 	fullyUploaded = true
@@ -572,6 +829,7 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 		// Cannot continue without metadata
 		return false, err
 	}
+	streamStats.MetaPointsCount = streamInfo.TotalPoints
 	streamStats.MetaPartsCount = streamInfo.PartsCount
 	if streamInfo.PartsCount == 0 {
 		// Edge device didn't finish uploading the stream yet
@@ -595,12 +853,17 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 			return false, objInfo.Err
 		}
 
+		if gAppQuitting {
+			logger.Infof("Quitting, stopping uploading one stream")
+			return false, nil
+		}
+
 		logger.Tracef("Checking minio file `%s`", objInfo.Key)
 
 		if strings.HasSuffix(objInfo.Key, "/") {
 			continue
 		}
-		partName := util.LastElem(strings.Split(objInfo.Key, "/"))
+		partName := filepath.Base(objInfo.Key)
 		if partName == "metadata.json" {
 			hasMetadata = true
 			continue
@@ -608,15 +871,15 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 
 		hasSomething = true
 
-		partObjUploadStats := StreamObjectUploadStatistics{
+		objStat := StreamObjectUploadStatistics{
 			StartTime: time.Now(),
 			PartName:  partName,
 		}
 
 		if part_already_uploaded(app, streamName, objInfo.Key) {
-			partObjUploadStats.EndTime = time.Now()
-			partObjUploadStats.UpState = "repeated"
-			streamStats.Objects[objInfo.Key] = partObjUploadStats
+			objStat.EndTime = time.Now()
+			objStat.UpState = "repeated"
+			streamStats.Objects[objInfo.Key] = objStat
 
 			logger.Infof("Part `%s` of stream `%s` is already uploaded", objInfo.Key, streamName)
 			continue
@@ -635,31 +898,45 @@ func upload_one_stream(app AppCtx, streamName string) (fullyUploaded bool, err e
 
 		err := upload_one_part(app, partInfo.StreamInfo, partInfo.StreamName, partInfo.PartName)
 		if err != nil {
-			partObjUploadStats.EndTime = time.Now()
-			partObjUploadStats.UpState = "fail"
-			partObjUploadStats.Msg = err.Error()
-			streamStats.Objects[objInfo.Key] = partObjUploadStats
+			objStat.EndTime = time.Now()
+			objStat.UpState = "fail"
+			objStat.Msg = err.Error()
 
 			logger.Warnf("Failed to upload part `%s` of stream `%s` (took %v): %v", partInfo.PartName, partInfo.StreamName,
-				partObjUploadStats.EndTime.Sub(partObjUploadStats.StartTime), err)
+				objStat.EndTime.Sub(objStat.StartTime), err)
 			fullyUploaded = false
-			continue
+		} else {
+			// Mark the part as uploaded
+			//err = app.db.Create(&PartUploadRecord{StreamName: partInfo.StreamName, PartName: partInfo.PartName}).Error
+			part := PartUploadRecord{StreamName: partInfo.StreamName, PartName: partInfo.PartName}
+			err = app.db.Where(part).FirstOrCreate(&PartUploadRecord{}).Error
+			if err != nil {
+				logger.Warnf("Failed to mark part `%s` of stream `%s` as uploaded: %v", partInfo.PartName, partInfo.StreamName, err)
+			}
+
+			objStat.EndTime = time.Now()
+			objStat.UpState = "ok"
+
+			logger.Infof("Uploaded part `%s` of stream `%s`, took %v", objInfo.Key, streamName,
+				objStat.EndTime.Sub(objStat.StartTime))
 		}
 
-		// Mark the part as uploaded
-		//err = app.db.Create(&PartUploadRecord{StreamName: partInfo.StreamName, PartName: partInfo.PartName}).Error
-		part := PartUploadRecord{StreamName: partInfo.StreamName, PartName: partInfo.PartName}
-		err = app.db.Where(part).FirstOrCreate(&PartUploadRecord{}).Error
+		streamStats.Objects[objInfo.Key] = objStat
+		partNum, err := util.ExtractNumberFromString(partName)
 		if err != nil {
-			logger.Warnf("Failed to mark part `%s` of stream `%s` as uploaded: %v", partInfo.PartName, partInfo.StreamName, err)
+			// Not a part file? Skip
+			continue
+		}
+		status := "success"
+		if objStat.UpState != "ok" {
+			status = "failed"
+		}
+		msg := smsg.MakePartInsertToStckStatusChangedMsg(gMachineID, objStat.StartTime.UnixNano(),
+			streamName, partNum, streamStats.MetaPointsCount, status, objStat.Msg)
+		err = PublishMessage(&msg)
+		if err != nil {
+			logger.Errorf("send part insert to stck status changed message error: %s", err)
 		}
-
-		partObjUploadStats.EndTime = time.Now()
-		partObjUploadStats.UpState = "ok"
-		streamStats.Objects[objInfo.Key] = partObjUploadStats
-
-		logger.Infof("Uploaded part `%s` of stream `%s`, took %v", objInfo.Key, streamName,
-			partObjUploadStats.EndTime.Sub(partObjUploadStats.StartTime))
 	}
 	if !hasMetadata {
 		logger.Warnf("Stream `%s` has no metadata file, will retry later", streamName)
@@ -707,14 +984,14 @@ func upload_one_part(app AppCtx, streamInfo *StreamMetadata, streamName string,
 		}
 
 		// Use regex to extract the part index from the part name
-		partIndex := 0
+		partIndex := int64(0)
 		{
 			re := regexp.MustCompile(`part_(\d+)\.zst`)
 			matches := re.FindStringSubmatch(partName)
 			if len(matches) != 2 {
 				return fmt.Errorf("failed to extract part index from part name `%s`", partName)
 			}
-			partIndex, err = strconv.Atoi(matches[1])
+			partIndex, err = strconv.ParseInt(matches[1], 10, 64)
 			if err != nil {
 				return fmt.Errorf("failed to convert part index `%s` to integer: %w", matches[1], err)
 			}
@@ -726,7 +1003,7 @@ func upload_one_part(app AppCtx, streamInfo *StreamMetadata, streamName string,
 
 			// Check if the part data size is correct
 			if streamInfo.PartsCount != 0 {
-				left := len(partData)
+				left := int64(len(partData))
 				if partIndex < streamInfo.PartsCount-1 {
 					right := streamInfo.PointsPerPart * 8
 					if left != right {

+ 1 - 0
minio-into-stck-installer/update.sh

@@ -21,6 +21,7 @@ cp -rf $config_dir /tmp/minio-into-stck-tmp
 ./install.sh
 echo "正在恢复配置文件..."
 rsync -a --delete /tmp/minio-into-stck-tmp/config/ "$config_dir"
+chmod 777 -R $config_dir
 rm -rf /tmp/minio-into-stck-tmp
 echo "更新完成!"
 

+ 2 - 1
package.sh

@@ -10,7 +10,8 @@ rm -rf $ROOT/tmp/*
 
 # ----- Start build -----
 rm $ROOT/minio-into-stck/minio-into-stck
-CGO_ENABLED=0 go build -o $ROOT/minio-into-stck/minio-into-stck $ROOT/
+BUILDTIME_STR=`date '+%Y-%m-%d %H:%M:%S'`
+CGO_ENABLED=0 go build -ldflags "-X \"main.buildtime=$BUILDTIME_STR\"" -o $ROOT/minio-into-stck/minio-into-stck $ROOT/
 
 cp -rf $ROOT/config $ROOT/minio-into-stck/
 # ----- End build -----

+ 1 - 0
stck-nsq-msg

@@ -0,0 +1 @@
+Subproject commit f37fc38c142e1593c3dddb45ae101817663a1370

+ 94 - 0
util/dchan.go

@@ -0,0 +1,94 @@
+package util
+
+import (
+	"sync"
+	"time"
+)
+
+// Deduplicated (debounced) channel. To close the channel, call `close(dchan.In())`.
+type DChan[T comparable] struct {
+	inChan        chan T
+	outChan       chan T
+	existingElems map[T]struct{}
+	mtx           sync.Mutex
+}
+
+func NewDChan[T comparable](capacity int) *DChan[T] {
+	dchan := &DChan[T]{
+		inChan:        make(chan T, capacity),
+		outChan:       make(chan T),
+		existingElems: make(map[T]struct{}),
+	}
+	go func() {
+		defer close(dchan.outChan)
+		for elem := range dchan.inChan {
+			dchan.mtx.Lock()
+			if _, ok := dchan.existingElems[elem]; !ok {
+				// Send the element to the output channel.
+				dchan.existingElems[elem] = struct{}{}
+				dchan.mtx.Unlock()
+				dchan.outChan <- elem
+			} else {
+				// Deduplicate the element.
+				dchan.mtx.Unlock()
+			}
+		}
+	}()
+	return dchan
+}
+
+func (dc *DChan[T]) In() chan<- T {
+	return dc.inChan
+}
+
+func (dc *DChan[T]) Out() <-chan T {
+	return dc.outChan
+}
+
+// Marks an element as read (stops deduplication), so that it can be sent and received again.
+func (dc *DChan[T]) MarkElemReadDone(elem T) {
+	dc.mtx.Lock()
+	delete(dc.existingElems, elem)
+	dc.mtx.Unlock()
+}
+
+// Checks if an element is in the channel.
+func (dc *DChan[T]) Contains(elem T) bool {
+	dc.mtx.Lock()
+	_, ok := dc.existingElems[elem]
+	dc.mtx.Unlock()
+	return ok
+}
+
+// Writes an element to the input channel.
+func (dc *DChan[T]) Write(elem T) {
+	dc.inChan <- elem
+}
+
+// Queues an element to be written to the input channel after a delay. If already in the channel, the delay is ignored.
+func (dc *DChan[T]) DelayedWrite(elem T, delay time.Duration) {
+	if delay <= 0 {
+		dc.inChan <- elem
+		return
+	}
+
+	// Check if the element is already in the channel.
+	dc.mtx.Lock()
+	if _, ok := dc.existingElems[elem]; ok {
+		// Existing element, abort.
+		dc.mtx.Unlock()
+		return
+	} else {
+		// New element, queue it.
+		dc.existingElems[elem] = struct{}{}
+		dc.mtx.Unlock()
+	}
+
+	// Write the element after the delay.
+	go func() {
+		time.Sleep(delay)
+		// HACK!
+		dc.MarkElemReadDone(elem)
+		dc.inChan <- elem
+	}()
+}

+ 72 - 0
util/file_util.go

@@ -0,0 +1,72 @@
+package util
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+)
+
+func CopyFile(src, dst string) error {
+	source, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer source.Close()
+
+	destDir := filepath.Dir(dst)
+	if err := os.MkdirAll(destDir, os.ModePerm); err != nil {
+		return err
+	}
+
+	destination, err := os.Create(dst)
+	if err != nil {
+		return err
+	}
+	defer destination.Close()
+
+	_, err = io.Copy(destination, source)
+	return err
+}
+
+func CreateHardLink(target, linkName string) error {
+	// 文件夹不存在则创建
+	destDir := filepath.Dir(linkName)
+	if err := os.MkdirAll(destDir, os.ModePerm); err != nil {
+		return err
+	}
+
+	return os.Link(target, linkName)
+}
+
+func FileExists(filename string) bool {
+	_, err := os.Stat(filename)
+	return !os.IsNotExist(err)
+}
+
+func RemoveEmptyDir(path string) error {
+	entries, err := os.ReadDir(path)
+	if err != nil {
+		return err
+	}
+
+	if len(entries) == 0 {
+		return os.Remove(path)
+	}
+
+	return nil
+}
+
+func EnsureFileCreated0777(path string) error {
+	if err := os.MkdirAll(filepath.Dir(path), 0777); err != nil {
+		return fmt.Errorf("create dir `%s` error: %v", filepath.Dir(path), err)
+	}
+	if _, err := os.OpenFile(path, os.O_CREATE, 0777); err != nil {
+		return fmt.Errorf("create file `%s` error: %v", path, err)
+	}
+	err := os.Chmod(path, 0777)
+	if err != nil {
+		return fmt.Errorf("chmod file `%s` error: %v", path, err)
+	}
+	return nil
+}

+ 48 - 0
util/util.go

@@ -1,11 +1,18 @@
 package util
 
 import (
+	"context"
+	"fmt"
+	"os"
 	"os/exec"
+	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
 	"github.com/klauspost/compress/zstd"
+	"github.com/minio/minio-go/v7"
 )
 
 func ExpandShellString(s string) (string, error) {
@@ -49,3 +56,44 @@ func FmtMyDate(t time.Time) string {
 func LastElem[T any](s []T) T {
 	return s[len(s)-1]
 }
+
+func MinioUploadFolder(minioClient *minio.Client, bucketName string, objPath string, localPath string) error {
+	// Walk the local folder
+	err := filepath.WalkDir(localPath, func(path string, d os.DirEntry, err error) error {
+		if err != nil {
+			return fmt.Errorf("walk dir failed: %w", err)
+		}
+
+		if d.IsDir() {
+			return nil
+		}
+
+		relPath, err := filepath.Rel(localPath, path)
+		if err != nil {
+			return fmt.Errorf("get relative path failed: %w", err)
+		}
+
+		targetPath := filepath.Join(objPath, relPath)
+		_, err = minioClient.FPutObject(context.Background(), bucketName, targetPath, path,
+			minio.PutObjectOptions{ContentType: "application/octet-stream"})
+		if err != nil {
+			return fmt.Errorf("fput object failed: %w", err)
+		}
+
+		return nil
+	})
+	return err
+}
+
+func GetCurrentNanoTimestampUTC() int64 {
+	return time.Now().UnixNano()
+}
+
+func ExtractNumberFromString(filename string) (int64, error) {
+	re := regexp.MustCompile(`(\d+)`)
+	match := re.FindStringSubmatch(filename)
+	if len(match) < 2 {
+		return 0, fmt.Errorf("number not found in filename `%s`", filename)
+	}
+	return strconv.ParseInt(match[1], 10, 64)
+}