From 44011db7258df050b74c89b05b6288acb8bef530 Mon Sep 17 00:00:00 2001
From: Vishal <vish9812@gmail.com>
Date: Thu, 15 Sep 2022 13:38:16 +0530
Subject: [PATCH] Add a semaphore (#21027)

---
 app/file.go   |  4 ++--
 app/server.go | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/app/file.go b/app/file.go
index a446a1d9d7..c461d1765b 100644
--- a/app/file.go
+++ b/app/file.go
@@ -682,7 +682,7 @@ func (a *App) UploadFileX(c *request.Context, channelID, name string, input io.R
 
 	if *a.Config().FileSettings.ExtractContent {
 		infoCopy := *t.fileinfo
-		a.Srv().Go(func() {
+		a.Srv().GoBuffered(func() {
 			err := a.ExtractContentFromFileInfo(&infoCopy)
 			if err != nil {
 				mlog.Error("Failed to extract file content", mlog.Err(err), mlog.String("fileInfoId", infoCopy.Id))
@@ -932,7 +932,7 @@ func (a *App) DoUploadFileExpectModification(c request.CTX, now time.Time, rawTe
 
 	if *a.Config().FileSettings.ExtractContent {
 		infoCopy := *info
-		a.Srv().Go(func() {
+		a.Srv().GoBuffered(func() {
 			err := a.ExtractContentFromFileInfo(&infoCopy)
 			if err != nil {
 				mlog.Error("Failed to extract file content", mlog.Err(err), mlog.String("fileInfoId", infoCopy.Id))
diff --git a/app/server.go b/app/server.go
index b44c19e9b2..3e3fd0a18b 100644
--- a/app/server.go
+++ b/app/server.go
@@ -134,6 +134,7 @@ type Server struct {
 
 	goroutineCount      int32
 	goroutineExitSignal chan struct{}
+	goroutineBuffered   chan struct{}
 
 	EmailService email.ServiceInterface
 
@@ -208,6 +209,7 @@ func NewServer(options ...Option) (*Server, error) {
 
 	s := &Server{
 		goroutineExitSignal: make(chan struct{}, 1),
+		goroutineBuffered:   make(chan struct{}, runtime.NumCPU()),
 		RootRouter:          rootRouter,
 		LocalRouter:         localRouter,
 		WebSocketRouter: &WebSocketRouter{
@@ -1037,6 +1039,26 @@ func (s *Server) Go(f func()) {
 	}()
 }
 
+// GoBuffered acts like a semaphore which creates a goroutine, but maintains a record of it
+// to ensure that execution completes before the server is shutdown.
+func (s *Server) GoBuffered(f func()) {
+	s.goroutineBuffered <- struct{}{}
+
+	atomic.AddInt32(&s.goroutineCount, 1)
+
+	go func() {
+		f()
+
+		atomic.AddInt32(&s.goroutineCount, -1)
+		select {
+		case s.goroutineExitSignal <- struct{}{}:
+		default:
+		}
+
+		<-s.goroutineBuffered
+	}()
+}
+
 // WaitForGoroutines blocks until all goroutines created by App.Go exit.
 func (s *Server) WaitForGoroutines() {
 	for atomic.LoadInt32(&s.goroutineCount) != 0 {