go语言实现并发遍历目录

go语言实现并发遍历目录

目录树结构定义

先实现一个目录树的数据结构定义，需要支持共享访问。

package main

import (
	"os"
	"strings"
	"sync"
)

type (
	// 目录项
	DirEntry struct {
		Name   string       `json:"name"`             // 文件（目录）名称
		Size   int64        `json:"size"`             // 大小
		IsDir  bool         `json:"is_dir"`           // 是否是目录
		Childs []*DirEntry  `json:"childs,omitempty"` // （如果是）目录下的目录项
		rwlock sync.RWMutex // 读写锁，用于多线程访问
	}
	// 目录树
	DirTree struct {
		Dir    string       `json:"dir"`  // 目录树起始路径
		Root   *DirEntry    `json:"root"` // 目录树起始路径节点
		rwlock sync.RWMutex // 读写锁，用于多线程访问
	}
)

// 获取目录下的子项
func (de *DirEntry) Get(name string) *DirEntry {
	de.rwlock.RLock()
	defer de.rwlock.RUnlock()

	for _, d := range de.Childs {
		if d.Name == name {
			return d
		}
	}
	return nil
}

// 添加目录项到目录
func (de *DirEntry) Add(child *DirEntry) {
	de.rwlock.Lock()
	defer de.rwlock.Unlock()

	for i := 0; i < len(de.Childs); i++ {
		if de.Childs[i].Name == child.Name {
			de.Childs[i] = child
			return
		}
	}
	de.Childs = append(de.Childs, child)
}

// 设置目录下的子项
func (de *DirEntry) SetChilds(childs []*DirEntry) {
	de.rwlock.Lock()
	defer de.rwlock.Unlock()

	de.Childs = childs
}

func (de *DirEntry) SetChildsFromFileinfo(childs []os.FileInfo) {
	de.rwlock.Lock()
	defer de.rwlock.Unlock()

	de.Childs = make([]*DirEntry, len(childs))

	for i, child := range childs {
		de.Childs[i] = NewDirEntry(child, nil)
	}
}

// 创建目录项（如果该目录项是文件，则childs传nil）
func NewDirEntry(d os.FileInfo, childs []os.FileInfo) *DirEntry {
	de := &DirEntry{
		Name:  d.Name(),
		Size:  d.Size(),
		IsDir: d.IsDir(),
	}
	if len(childs) == 0 {
		return de
	}

	de.Childs = make([]*DirEntry, len(childs))

	for i, child := range childs {
		de.Childs[i] = NewDirEntry(child, nil)
	}
	return de
}

// 获取目录树下的目录项
//
// path 要获取的目录项路径
func (dt *DirTree) Get(path string) *DirEntry {
	if !strings.HasPrefix(path, dt.Dir) {
		return nil
	}
	path = path[len(dt.Dir):]

	pathparts := strings.FieldsFunc(path, func(r rune) bool { return r == '/' })

	dt.rwlock.RLock()
	defer dt.rwlock.RUnlock()

	de := dt.Root
	for _, part := range pathparts {
		if de == nil {
			break
		}
		de = de.Get(part)
	}

	return de
}

并发遍历目录实现代码

使用go语言实现并发访问是很容易的，go的协程机制很方便实现。

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Printf("使用:%s dir(or file)\n", os.Args[0])
		return
	}
	dir := os.Args[1]

	dt, err := pWalkDir(dir)
	if err != nil {
		fmt.Println("错误:", err)
		return
	}
	data, err := json.MarshalIndent(dt, "", "  ")
	if err != nil {
		fmt.Println("错误:", err)
		return
	}

	fmt.Println(string(data))

	fmt.Println("目录遍历完成")
}

// 并发遍历目录
func pWalkDir(dir string) (*DirTree, error) {

	fi, err := os.Stat(dir)
	if err != nil {
		return nil, err
	}
	if !fi.IsDir() {
		return nil, fmt.Errorf("不是目录")
	}
	// 创建目录树
	dt := &DirTree{
		Dir:  dir,
		Root: NewDirEntry(fi, nil),
	}
	fmt.Println(dt.Root.Name)

	// 设置并发数限制
	concurrencyLimit := make(chan struct{}, 4)

	// 递归遍历目录
	var wg sync.WaitGroup
	wg.Add(1)
	pWalkDirImpl(dir, dt, concurrencyLimit, &wg)
	wg.Wait()

	return dt, nil
}

// 实现并发递归遍历目录
func pWalkDirImpl(dir string, dt *DirTree, concurrencyLimit chan struct{}, wg *sync.WaitGroup) error {
	defer wg.Done()

	// 读取目录
	fis, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}
	// 获取子项中的目录列表
	dirs := make([]string, 0, 16)
	for _, fi := range fis {
		path := filepath.Join(dir, fi.Name())
		if fi.IsDir() {
			dirs = append(dirs, path)
			continue
		}
	}
	// 更新数据到目录树
	de := dt.Get(dir)
	if de == nil {
		fmt.Println(dir, dt.Dir, dt.Root.Name)
		os.Exit(0)
	}
	de.SetChildsFromFileinfo(fis)

	// 子项中目录再遍历
	for _, d := range dirs {
		concurrencyLimit <- struct{}{}
		wg.Add(1)
		go pWalkDirImpl(d, dt, concurrencyLimit, wg)
		<-concurrencyLimit
	}
	return nil
}

测试输出

上面输出是使用json格式，也可以根据需要改成类似tree输出。

改成树形输出，使用[github.com/disiqueira/gotree](github.com/disiqueira/gotree)库，实现Tree和Printer接口即可。

mongodb
{
  "dir": "../../mongodb",
  "root": {
    "name": "mongodb",
    "size": 4096,
    "is_dir": true,
    "childs": [
      {
        "name": "mongodb",
        "size": 4096,
        "is_dir": true,
        "childs": [
          {
            "name": "mongodb-org-server-6.0.5-1.el8.x86_64.rpm",
            "size": 32301236,
            "is_dir": false
          },
          {
            "name": "mongodb-org-tools-6.0.5-1.el8.x86_64.rpm",
            "size": 10788,
            "is_dir": false
          }
        ]
      },
      {
        "name": "percona-server-mongodb-6.0.5-4.el8.x86_64",
        "size": 4096,
        "is_dir": true,
        "childs": [
          {
            "name": "percona-mongodb-mongosh-1.8.0.el8.x86_64.rpm",
            "size": 25670076,
            "is_dir": false
          },
          {
            "name": "percona-server-mongodb-6.0.5-4.el8.x86_64.rpm",
            "size": 8352,
            "is_dir": false
          }
        ]
      }
    ]
  }
}