├── .codecov.yml ├── .github └── workflows │ ├── codeql-analysis.yml │ └── go.yml ├── .gitignore ├── CHANGELOG.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── README_CN.md ├── docs ├── COPYRIGHT.txt ├── images │ └── waterdrop.jpg └── zh-cn │ ├── distribute-lock.md │ ├── install-dependencies.md │ ├── quick-start.md │ └── tool.md ├── go.mod ├── go.sum ├── pkg ├── breaker │ ├── breaker.go │ ├── googlesrebreaker.go │ └── googlesrebreaker_test.go ├── broker │ ├── kafka │ │ ├── README.md │ │ ├── consumer.go │ │ ├── examples │ │ │ ├── consumer │ │ │ │ └── consumer.go │ │ │ └── producer │ │ │ │ └── producer.go │ │ ├── interceptor.go │ │ └── producer.go │ └── rocketmq │ │ ├── README.md │ │ ├── consumer.go │ │ ├── examples │ │ ├── consumer │ │ │ └── consumer.go │ │ └── producer │ │ │ └── producer.go │ │ ├── interceptor.go │ │ ├── primitive.go │ │ └── producer.go ├── conf │ ├── conf.go │ ├── parser │ │ └── toml │ │ │ └── toml.go │ └── provider │ │ └── file │ │ └── file.go ├── database │ ├── es │ │ ├── client.go │ │ ├── es_test.go │ │ ├── query.go │ │ └── stats.go │ ├── mongo │ │ ├── aggregate.go │ │ ├── bulk.go │ │ ├── client.go │ │ ├── collection.go │ │ ├── collection_test.go │ │ └── query.go │ ├── redis │ │ ├── command.go │ │ ├── lock.go │ │ ├── lock_test.go │ │ ├── redis.go │ │ └── redis_test.go │ └── sql │ │ ├── mysql.go │ │ ├── mysql_test.go │ │ ├── postgres.go │ │ └── sql.go ├── log │ ├── log.go │ └── log_test.go ├── ratelimit │ └── setinel │ │ ├── sentinel.go │ │ └── sentinel_test.go ├── registry │ ├── etcd │ │ ├── etcd.go │ │ ├── etcd_test.go │ │ ├── lock.go │ │ └── lock_test.go │ └── registry.go ├── server │ ├── http │ │ ├── client │ │ │ ├── client.go │ │ │ └── client_test.go │ │ ├── config │ │ │ └── config.go │ │ ├── metadata │ │ │ └── metadata.go │ │ ├── middlewares │ │ │ ├── cors.go │ │ │ ├── cors_test.go │ │ │ ├── logger.go │ │ │ ├── logger_test.go │ │ │ ├── metric.go │ │ │ ├── metric_test.go │ │ │ ├── ratelimit │ │ │ │ ├── limiter.go │ │ │ │ └── sentinel │ │ │ │ │ ├── sentinel.go │ │ │ │ │ └── sentinel_test.go │ │ │ ├── recovery.go │ │ │ ├── recovery_test.go │ │ │ ├── trace.go │ │ │ └── trace_test.go │ │ ├── server │ │ │ ├── server.go │ │ │ └── server_test.go │ │ └── websocket │ │ │ └── websocket.go │ └── rpc │ │ ├── client │ │ ├── client.go │ │ └── client_test.go │ │ ├── config │ │ └── config.go │ │ ├── interceptors │ │ ├── breaker.go │ │ ├── breaker_test.go │ │ ├── logger.go │ │ ├── logger_test.go │ │ ├── metric.go │ │ ├── metric_test.go │ │ ├── recovery.go │ │ ├── recovery_test.go │ │ ├── sentinel.go │ │ ├── sentinel_test.go │ │ ├── trace.go │ │ ├── trace_test.go │ │ ├── validator.go │ │ └── validator_test.go │ │ ├── metadata │ │ └── metadata.go │ │ └── server │ │ ├── server.go │ │ └── server_test.go ├── stats │ ├── metric │ │ ├── const.go │ │ ├── counter.go │ │ ├── gauge.go │ │ ├── histogram.go │ │ └── metric.go │ ├── profile │ │ └── profile.go │ └── stats.go ├── status │ ├── README.md │ └── status.go ├── store │ ├── metadata.go │ └── minio │ │ ├── minio.go │ │ └── minio_test.go ├── trace │ ├── jaeger │ │ └── jaeger.go │ ├── metadata.go │ └── trace.go ├── utils │ ├── xbuffer │ │ ├── README.md │ │ ├── bufferpool.go │ │ ├── bufferpool_test.go │ │ ├── sizedbufferpool.go │ │ └── sizedbufferpool_test.go │ ├── xcollection │ │ ├── lru.go │ │ ├── lru_test.go │ │ ├── priorityqueue.go │ │ ├── priorityqueue_test.go │ │ ├── rollingwindow.go │ │ ├── rollingwindow_test.go │ │ ├── safemap.go │ │ ├── safemap_test.go │ │ ├── trie.go │ │ └── trie_test.go │ ├── xcrypto │ │ ├── aes.go │ │ ├── aes_test.go │ │ ├── hash.go │ │ ├── hash_test.go │ │ ├── pem │ │ │ ├── rsa_private_key.pem │ │ │ └── rsa_public_key.pem │ │ ├── rsa.go │ │ ├── rsa_test.go │ │ └── util.go │ ├── xdefer │ │ ├── defers.go │ │ └── defers_test.go │ ├── xnet │ │ ├── ip.go │ │ └── ip_test.go │ ├── xslice │ │ ├── float.go │ │ ├── float_test.go │ │ ├── int.go │ │ ├── int_test.go │ │ ├── string.go │ │ └── string_test.go │ ├── xstring │ │ ├── string.go │ │ └── string_test.go │ ├── xsync │ │ ├── singlecall.go │ │ ├── singlecall_test.go │ │ ├── waitgroup.go │ │ └── waitgroup_test.go │ └── xtime │ │ ├── jitter.go │ │ ├── jitter_test.go │ │ ├── time.go │ │ └── time_test.go └── version │ └── version.go ├── tests └── proto │ └── demo │ ├── demo.pb.go │ └── demo.proto └── tools └── waterdrop ├── ecode └── i18n.go ├── go.mod ├── go.sum ├── main.go ├── project ├── README.md ├── base │ ├── install.go │ ├── install_compatible.go │ ├── mod.go │ ├── mod_test.go │ ├── path.go │ ├── repo.go │ └── repo_test.go └── layout.go ├── protoc ├── action.go ├── cmd.go ├── common.go ├── grpc.go └── swagger.go ├── swagger └── cmd.go ├── testgen ├── common │ ├── parser.go │ └── utils.go └── utgen │ ├── cmd.go │ ├── gen.go │ └── templete.go ├── upgrade └── upgrade.go └── utils ├── env.go └── helper.go /.codecov.yml: -------------------------------------------------------------------------------- 1 | ignore: 2 | - "docs" 3 | - "tests" 4 | - "tools" -------------------------------------------------------------------------------- /.github/workflows/codeql-analysis.yml: -------------------------------------------------------------------------------- 1 | # For most projects, this workflow file will not need changing; you simply need 2 | # to commit it to your repository. 3 | # 4 | # You may wish to alter this file to override the set of languages analyzed, 5 | # or to provide custom queries or build logic. 6 | # 7 | # ******** NOTE ******** 8 | # We have attempted to detect the languages in your repository. Please check 9 | # the `language` matrix defined below to confirm you have the correct set of 10 | # supported CodeQL languages. 11 | # 12 | name: "CodeQL" 13 | 14 | on: 15 | push: 16 | branches: [ master ] 17 | pull_request: 18 | # The branches below must be a subset of the branches above 19 | branches: [ master ] 20 | schedule: 21 | - cron: '35 15 * * 4' 22 | 23 | jobs: 24 | analyze: 25 | name: Analyze 26 | runs-on: ubuntu-latest 27 | permissions: 28 | actions: read 29 | contents: read 30 | security-events: write 31 | 32 | strategy: 33 | fail-fast: false 34 | matrix: 35 | language: [ 'go' ] 36 | # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ] 37 | # Learn more: 38 | # https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed 39 | 40 | steps: 41 | - name: Checkout repository 42 | uses: actions/checkout@v2 43 | 44 | # Initializes the CodeQL tools for scanning. 45 | - name: Initialize CodeQL 46 | uses: github/codeql-action/init@v1 47 | with: 48 | languages: ${{ matrix.language }} 49 | # If you wish to specify custom queries, you can do so here or in a config file. 50 | # By default, queries listed here will override any specified in a config file. 51 | # Prefix the list here with "+" to use these queries and those in the config file. 52 | # queries: ./path/to/local/query, your-org/your-repo/queries@main 53 | 54 | # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). 55 | # If this step fails, then you should remove it and run the build manually (see below) 56 | - name: Autobuild 57 | uses: github/codeql-action/autobuild@v1 58 | 59 | # ℹ️ Command-line programs to run using the OS shell. 60 | # 📚 https://git.io/JvXDl 61 | 62 | # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines 63 | # and modify them (or add more) to build your code if your project 64 | # uses a compiled language 65 | 66 | #- run: | 67 | # make bootstrap 68 | # make release 69 | 70 | - name: Perform CodeQL Analysis 71 | uses: github/codeql-action/analyze@v1 72 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Reference https://github.com/github/gitignore/blob/master/Go.gitignore 2 | # Binaries for programs and plugins 3 | *.exe 4 | *.exe~ 5 | *.dll 6 | *.so 7 | *.dylib 8 | 9 | # Test binary, built with `go test -c` 10 | *.test 11 | 12 | # Output of the go coverage tool, specifically when used with LiteIDE 13 | *.out 14 | 15 | # Dependency directories (remove the comment below to include it) 16 | vendor/ 17 | 18 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 19 | *.o 20 | *.a 21 | *.so 22 | 23 | # OS General 24 | Thumbs.db 25 | .DS_Store 26 | 27 | # project 28 | *.cert 29 | *.key 30 | *.log 31 | bin/ 32 | 33 | # Develop tools 34 | .vscode/ 35 | .idea/ 36 | *.swp 37 | 38 | # project directories and files 39 | logs/ 40 | coverage.txt 41 | 42 | #binary package 43 | tools/waterdrop/main 44 | tools/waterdrop/waterdrop 45 | tools/waterdrop/*.proto 46 | tools/waterdrop/*.pb.go 47 | tools/waterdrop/*.json -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contribute to Waterdrop 2 | 3 | The following is a set of guidelines for contributing to waterdrop. If you find something incorrect or missing, please leave comments, suggestions or file bug. 4 | 5 | ## Reporting issues 6 | 7 | Reporting issues are a great way to contribute to the project. We are perpetually grateful about a well-written, thorough bug report. 8 | 9 | Before raising a new issue, check our [issue list](https://github.com/UnderTreeTech/waterdrop/issues) to determine if it already contains the problem that you are facing. 10 | 11 | A good bug report shouldn't leave others needing to chase you for more information. Please be as detailed as possible. The following questions might serve as a template for writing a detailed report: 12 | 13 | - What were you trying to achieve? 14 | - What are the expected results? 15 | - What are the received results? 16 | - What are the steps to reproduce the issue? 17 | - In what environment did you encounter the issue? 18 | 19 | Feature requests can also be submitted as issues. 20 | 21 | 22 | ## Pull requests 23 | 24 | Good pull requests (e.g. patches, improvements, new features) are a fantastic help. They should remain focused in scope and avoid unrelated commits. 25 | 26 | Please ask first before embarking on any significant pull request (e.g. implementing new features, refactoring code etc.), otherwise you risk spending a lot of time working on something that the maintainers might not want to merge into the project. 27 | 28 | First add an issue to the project to discuss the improvement. Please adhere to the coding conventions used throughout the project. If in doubt, consult the [Effective Go style guide](https://golang.org/doc/effective_go.html) and [Uber Go style guide](https://github.com/uber-go/guide). 29 | -------------------------------------------------------------------------------- /README_CN.md: -------------------------------------------------------------------------------- 1 | 简体中文 | [English](README.md) 2 | 3 | [![Go](https://github.com/UnderTreeTech/waterdrop/workflows/Go/badge.svg?branch=master)](https://github.com/UnderTreeTech/waterdrop/actions) 4 | [![codecov](https://codecov.io/gh/UnderTreeTech/waterdrop/branch/master/graph/badge.svg)](https://codecov.io/gh/UnderTreeTech/waterdrop) 5 | [![Go Report Card](https://goreportcard.com/badge/github.com/UnderTreeTech/waterdrop)](https://goreportcard.com/report/github.com/UnderTreeTech/waterdrop) 6 | [![Release](https://img.shields.io/github/v/release/UnderTreeTech/waterdrop.svg?style=flat-square)](https://github.com/UnderTreeTech/waterdrop) 7 | [![License](https://img.shields.io/badge/license-Apache%202-4EB1BA.svg)](https://www.apache.org/licenses/LICENSE-2.0.html) 8 | 9 | ## Waterdrop 10 | 11 | 水滴是基于gin及grpc构建的一款高性能微服务框架. 水滴之名来自于《三体》——纯洁而唯美，攻击高速有效，威力巨大！ 12 | ![水滴](docs/images/waterdrop.jpg) 13 | 14 | ## Features 15 | 16 | - HTTP Server：基于gin进行封装，可复用gin所有特性 17 | - RPC Server：基于官方gRPC开发，基于etcd进行服务注册发现，默认roundrobin负载均衡 18 | - Conf：支持yaml、toml、json等多格式扩展，默认toml解析，可自定义是否监听文件热更新配置 19 | - Database：集成MySQL、Redis 20 | - Log：基于Zap封装 21 | - Trace：集成Opentracing接入，jaeger落地支撑 22 | - Distribute Lock：基于Redis、ETCD实现分布式锁，前者适合最终一致性业务锁，后者适合强一致性业务锁 23 | - Stats：服务运行metrics & profile 24 | - Broker：默认支持RocketMQ, Kafka. 25 | - Utils: 辅助类函数 26 | - Registry：服务注册发现，制定通用接口定义，默认支持etcd 27 | - Status：全局错误处理，用于HTTP/RPC之间错误转换。后续可扩展成从remote加载错误定义 28 | - Dashboard：基于Grafana搭建metrics大盘，待实现 29 | - Breaker：熔断器，支持[alibaba sentinel](https://github.com/alibaba/sentinel-golang)、[google sre breaker](https://landing.google.com/sre/sre-book/chapters/handling-overload/) 30 | - Middlewares & Interceptors：http/rpc server通用中间件，如recovery, trace, metric and logger等 31 | 32 | 33 | ## Installation 34 | 35 | `go get github.com/UnderTreeTech/waterdrop` 36 | 37 | ## Tools 38 | 39 | waterdrop提供脚手架工具来提高开发效率。执行`go get -u github.com/UnderTreeTech/waterdrop/tools/waterdrop`安装最新版工具。 40 | 41 | 工具依赖`protc`及`protoc-gen-go`来生成protobuf代码，目前waterdrop工具并不自动安装这两个插件需要用户自主安装，实际开发中每人的版本并不相同。 42 | 43 | waterdrop工具提供的功能如下： 44 | 45 | - `waterdrop new your_project_name` new a standard layout project 46 | - `waterdrop protoc --grpc --swagger xx.proto` generate grpc code and swagger api file 47 | - `waterdrop swagger serve xx.swagger.json` serve and browse swagger api 48 | - `waterdrop utgen xx.go` generate unit tests 49 | - `waterdrop upgrade` upgrade tool `waterdrop` 50 | 51 | ## Contributing 52 | 53 | Contributions are always welcomed! You can start with the issues labeled with bug or feature. 54 | 55 | -------------------------------------------------------------------------------- /docs/COPYRIGHT.txt: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2024 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | -------------------------------------------------------------------------------- /docs/images/waterdrop.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/UnderTreeTech/waterdrop/87ded0949f353217362115fdcc19ecb53d372dd7/docs/images/waterdrop.jpg -------------------------------------------------------------------------------- /docs/zh-cn/distribute-lock.md: -------------------------------------------------------------------------------- 1 | ## Distribute Lock 2 | 3 | ### Redis Lock 4 | 适合最终一致的业务锁，代码路径database/redis/lock.go 5 | ```go 6 | Lock 7 | UnLock 8 | ForceUnlock 9 | ``` 10 | 11 | ### Etcd lock 12 | 适合强一致的业务锁，代码路径registry/etcd/lock.go 13 | ```go 14 | NewMutex 15 | Lock 16 | TryLock 17 | UnLock 18 | ``` -------------------------------------------------------------------------------- /docs/zh-cn/install-dependencies.md: -------------------------------------------------------------------------------- 1 | # 依赖安装 2 | 3 | ## etcd 4 | 5 | 官方 [etcd](https://github.com/etcd-io/etcd/releases) 下载二进制包解压，进入解压后目录并执行./etcd即可，也可以将可执行文件加入到全局环境变量中 6 | 7 | ## jaeger 8 | 9 | 官方 [jaeger](https://github.com/jaegertracing/jaeger/releases) 下载二进制包解压，进入解压后目录并执行./jaeger-all-in-one即可，也可以将可执行文件加入到全局环境变量中 10 | 11 | ## redis 12 | 13 | 1、浏览器中下载 [redis-5.0.9](https://download.redis.io/releases/redis-5.0.9.tar.gz) 14 | 15 | 2、执行如下命令 16 | ```shell 17 | tar xzf redis-5.0.9.tar.gz 18 | cd redis-5.0.9 19 | make 20 | 21 | ./src/redis-server 22 | 23 | ``` 24 | 25 | ## mysql 26 | 27 | 官方下载 [mysql](https://dev.mysql.com/downloads/mysql/) 安装包安装 28 | 29 | -------------------------------------------------------------------------------- /docs/zh-cn/quick-start.md: -------------------------------------------------------------------------------- 1 | ## Quick Start 2 | 3 | 1、 安装依赖 4 | 5 | 安装 [dependencies](install-dependencies.md) 6 | 7 | 2、 拉取项目 8 | 9 | `go get -u github.com/UnderTreeTech/waterdrop` 10 | 11 | 3、 进入examples运行示例 12 | 13 | ## 生成模板——待实现 14 | 15 | `waterdrop new your_project_name` 生成项目模板 -------------------------------------------------------------------------------- /docs/zh-cn/tool.md: -------------------------------------------------------------------------------- 1 | # waterdrop工具 2 | 3 | Clone [waterdrop](https://github.com/UnderTreeTech/waterdrop) 代码，进入tools目录，执行go install 4 | 5 | ## 依赖安装 6 | 7 | 安装protoc：官方下载 [protobuf](https://github.com/protocolbuffers/protobuf/releases) 对应平台的二进制包，解压后按照README进行安装即可 8 | 9 | Notice：waterdrop采用 [gogo protobuf](https://github.com/golang/protobuf) 做为pb生成工具 10 | 11 | 安装protoc-gen-go：执行命令 go get -u github.com/golang/protobuf/protoc-gen-go 12 | 13 | ## pb代码生成器 14 | 15 | 执行 `waterdrop protoc --grpc your.proto` 即可生成pb代码 16 | 17 | ## swagger接口文档 18 | 19 | 执行 `waterdrop protoc --swagger your.proto` 即可生成对应的swagger api文档 20 | 21 | 执行 `waterdrop swagger serve your.swagger.json` 可即时预览api定义 22 | 23 | ## 单元测试UT 24 | 25 | 执行 `waterdrop utgen your.go` 即可生成该文件下所有方法的单元测试 26 | 执行 `waterdrop utgen --func XXXX your.go` 即可生成对应文件下某个方法的单元测试 27 | 28 | ## 说明 29 | 30 | 以上操作第一次执行时会相对比较耗时，需要下载各工具相应的依赖包 -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/UnderTreeTech/waterdrop 2 | 3 | go 1.14 4 | 5 | require ( 6 | github.com/BurntSushi/toml v1.1.0 7 | github.com/Shopify/sarama v1.32.0 8 | github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect 9 | github.com/alibaba/sentinel-golang v1.0.2 10 | github.com/apache/rocketmq-client-go/v2 v2.1.2 11 | github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd // indirect 12 | github.com/gin-contrib/cors v1.3.1 13 | github.com/gin-gonic/gin v1.7.7 14 | github.com/go-ole/go-ole v1.2.4 // indirect 15 | github.com/go-playground/validator/v10 v10.4.1 16 | github.com/go-redis/redis/v8 v8.11.5 17 | github.com/go-resty/resty/v2 v2.7.0 18 | github.com/go-sql-driver/mysql v1.6.0 19 | github.com/golang/mock v1.4.4 // indirect 20 | github.com/google/uuid v1.3.0 21 | github.com/gorilla/websocket v1.4.2 22 | github.com/json-iterator/go v1.1.12 23 | github.com/lib/pq v1.8.0 24 | github.com/minio/minio-go/v7 v7.0.12 25 | github.com/mitchellh/mapstructure v1.4.1 26 | github.com/natefinch/lumberjack v0.0.0-20201021141957-47ffae23317c 27 | github.com/olivere/elastic/v7 v7.0.25 28 | github.com/opentracing/opentracing-go v1.2.0 29 | github.com/prometheus/client_golang v1.11.1 30 | github.com/qiniu/qmgo v1.0.6 31 | github.com/spf13/cast v1.4.1 32 | github.com/stretchr/objx v0.2.0 // indirect 33 | github.com/stretchr/testify v1.7.0 34 | github.com/uber/jaeger-client-go v2.30.0+incompatible 35 | github.com/uber/jaeger-lib v2.2.0+incompatible // indirect 36 | go.etcd.io/etcd/api/v3 v3.5.0 37 | go.etcd.io/etcd/client/v3 v3.5.0 38 | go.mongodb.org/mongo-driver v1.8.2 39 | go.uber.org/automaxprocs v1.4.0 40 | go.uber.org/zap v1.21.0 41 | google.golang.org/genproto v0.0.0-20210602131652-f16073e35f0c 42 | google.golang.org/grpc v1.40.0 43 | google.golang.org/protobuf v1.26.0 44 | gopkg.in/fsnotify.v1 v1.4.7 45 | ) 46 | -------------------------------------------------------------------------------- /pkg/breaker/breaker.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package breaker 20 | 21 | import ( 22 | "math/rand" 23 | "sync" 24 | "time" 25 | ) 26 | 27 | const ( 28 | // StateClosed when circuit breaker closed, request allowed, the breaker 29 | // calc the succeed ratio, if request num greater request setting and 30 | // ratio lower than the setting ratio, then reset state to open. 31 | StateClosed int32 = iota 32 | // StateOpen when circuit breaker open, request not allowed, after sleep 33 | // some duration, allow one single request for testing the health, if ok 34 | // then state reset to closed, if not continue the step. 35 | StateOpen 36 | ) 37 | 38 | type Breaker interface { 39 | Allow() error 40 | Accept() 41 | Reject() 42 | } 43 | 44 | type Proba struct { 45 | r *rand.Rand 46 | lock sync.Mutex 47 | } 48 | 49 | // NewProba return Proba pointer 50 | func NewProba() *Proba { 51 | return &Proba{ 52 | r: rand.New(rand.NewSource(time.Now().UnixNano())), 53 | } 54 | } 55 | 56 | // TrueOnProba check if input proba less than pseudo-random number 57 | func (p *Proba) TrueOnProba(proba float64) bool { 58 | p.lock.Lock() 59 | reject := p.r.Float64() < proba 60 | p.lock.Unlock() 61 | return reject 62 | } 63 | 64 | // brks global breaker group instance 65 | var brks *BreakerGroup 66 | 67 | func init() { 68 | brks = &BreakerGroup{ 69 | breakers: make(map[string]Breaker), 70 | } 71 | } 72 | 73 | // BreakerGroup brks 74 | type BreakerGroup struct { 75 | mutex sync.RWMutex 76 | breakers map[string]Breaker 77 | } 78 | 79 | // NewBreakerGroup returns global breaker group instance brks 80 | func NewBreakerGroup() *BreakerGroup { 81 | return brks 82 | } 83 | 84 | // Get return a break associate with the name 85 | func (bg *BreakerGroup) Get(name string) Breaker { 86 | bg.mutex.RLock() 87 | breaker, ok := bg.breakers[name] 88 | bg.mutex.RUnlock() 89 | if ok { 90 | return breaker 91 | } 92 | 93 | bg.mutex.Lock() 94 | breaker, ok = bg.breakers[name] 95 | if !ok { 96 | cfg := defaultGoogleSreBreakerConfig() 97 | cfg.Name = name 98 | breaker = newGoogleSreBreaker(cfg) 99 | bg.breakers[name] = breaker 100 | } 101 | bg.mutex.Unlock() 102 | 103 | return breaker 104 | } 105 | 106 | // Do execute the input func and stats the breaker result 107 | func (bg *BreakerGroup) Do(name string, run func() error, accept func(error) bool) error { 108 | breaker := bg.Get(name) 109 | err := breaker.Allow() 110 | if err != nil { 111 | return err 112 | } 113 | 114 | err = run() 115 | if accept(err) { 116 | breaker.Accept() 117 | } else { 118 | breaker.Reject() 119 | } 120 | 121 | return err 122 | } 123 | -------------------------------------------------------------------------------- /pkg/breaker/googlesrebreaker_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package breaker 20 | 21 | import ( 22 | "os" 23 | "testing" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/log" 26 | 27 | "github.com/stretchr/testify/assert" 28 | ) 29 | 30 | var bg *BreakerGroup 31 | 32 | func TestMain(m *testing.M) { 33 | defer log.New(nil).Sync() 34 | bg = NewBreakerGroup() 35 | os.Exit(m.Run()) 36 | } 37 | 38 | // TestBreakerAccept test breaker Accept 39 | func TestBreakerAccept(t *testing.T) { 40 | breaker := bg.Get("breaker") 41 | for i := 0; i < 100; i++ { 42 | breaker.Accept() 43 | } 44 | assert.Nil(t, breaker.Allow()) 45 | } 46 | 47 | // TestBreakerReject test breaker Reject 48 | func TestBreakerReject(t *testing.T) { 49 | breaker := bg.Get("breaker") 50 | for i := 0; i < 40000; i++ { 51 | breaker.Reject() 52 | } 53 | err := breaker.Allow() 54 | assert.NotNil(t, err) 55 | } 56 | 57 | // TestBreakerDo test breaker Do 58 | func TestBreakerDo(t *testing.T) { 59 | err := bg.Do("do", func() error { 60 | return nil 61 | }, func(e error) bool { 62 | return e == nil 63 | }) 64 | assert.Nil(t, err) 65 | assert.Panics(t, func() { 66 | bg.Do("do", func() error { 67 | panic("exit") 68 | }, func(e error) bool { 69 | return e == nil 70 | }) 71 | }) 72 | } 73 | -------------------------------------------------------------------------------- /pkg/broker/kafka/README.md: -------------------------------------------------------------------------------- 1 | ## Kafka 2 | 3 | Broker: Kafka 4 | 5 | You can run examples by replace configs use your kafka instance configs. -------------------------------------------------------------------------------- /pkg/broker/kafka/examples/consumer/consumer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package main 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "time" 25 | 26 | "github.com/Shopify/sarama" 27 | 28 | "github.com/UnderTreeTech/waterdrop/pkg/broker/kafka" 29 | "github.com/UnderTreeTech/waterdrop/pkg/log" 30 | ) 31 | 32 | func main() { 33 | defer log.New(nil).Sync() 34 | 35 | config := &kafka.ConsumerConfig{ 36 | Addr: []string{"your_instance_addr"}, 37 | Topic: []string{"your_instance_topic"}, 38 | Gid: "your_group_id", 39 | EnableSASLAuth: true, 40 | SASLMechanism: "your_sasl_mechanics", //PLAIN 41 | SASLUser: "your_sasl_user", 42 | SASLPassword: "your_sasl_password", 43 | SASLHandshake: true, 44 | 45 | DialTimeout: time.Second * 5, 46 | ConsumeOldest: true, 47 | EnableReturnError: true, 48 | } 49 | 50 | consumer := kafka.NewConsumer(config) 51 | consumer.Subscribe(consume) 52 | consumer.Start() 53 | 54 | time.Sleep(time.Hour) 55 | consumer.Close() 56 | } 57 | 58 | func consume(ctx context.Context, message *sarama.ConsumerMessage) error { 59 | fmt.Println(fmt.Sprintf("Message claimed: value = %s, timestamp = %v, topic = %s", string(message.Value), message.Timestamp, message.Topic)) 60 | return nil 61 | } 62 | -------------------------------------------------------------------------------- /pkg/broker/kafka/examples/producer/producer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package main 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "time" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/utils/xstring" 27 | 28 | "github.com/UnderTreeTech/waterdrop/pkg/broker/kafka" 29 | "github.com/UnderTreeTech/waterdrop/pkg/log" 30 | ) 31 | 32 | func main() { 33 | defer log.New(nil).Sync() 34 | 35 | config := &kafka.ProducerConfig{ 36 | Addr: []string{"your_instance_addr"}, 37 | Topic: []string{"your_instance_topic"}, 38 | EnableSASLAuth: true, 39 | SASLMechanism: "your_sasl_mechanics", //PLAIN 40 | SASLUser: "your_sasl_user", 41 | SASLPassword: "your_sasl_password", 42 | SASLHandshake: true, 43 | 44 | DialTimeout: time.Second * 5, 45 | EnableReturnSuccess: true, 46 | } 47 | 48 | producer := kafka.NewSyncProducer(config) 49 | 50 | for i := 0; i < 100000; i++ { 51 | err := producer.SendSyncMsg(context.Background(), xstring.RandomString(16)) 52 | if err != nil { 53 | fmt.Println("error", err.Error()) 54 | } 55 | } 56 | 57 | producer.Close() 58 | } 59 | -------------------------------------------------------------------------------- /pkg/broker/kafka/interceptor.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package kafka 20 | 21 | import "github.com/Shopify/sarama" 22 | 23 | type TraceInterceptor struct { 24 | TraceID string 25 | } 26 | 27 | func (ti *TraceInterceptor) OnSend(message *sarama.ProducerMessage) { 28 | 29 | } 30 | -------------------------------------------------------------------------------- /pkg/broker/rocketmq/README.md: -------------------------------------------------------------------------------- 1 | ## RocketMQ 2 | 3 | Broker: RocketMQ 4 | 5 | You can run examples by replace configs use your rocketmq instance configs. 6 | 7 | Typically log msg in interceptors, default provide producer and push consumer inteceptor. -------------------------------------------------------------------------------- /pkg/broker/rocketmq/examples/consumer/consumer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package main 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "time" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/broker/rocketmq" 27 | "github.com/UnderTreeTech/waterdrop/pkg/log" 28 | ) 29 | 30 | func main() { 31 | defer log.New(nil).Sync() 32 | 33 | config := &rocketmq.ConsumerConfig{ 34 | Endpoint: []string{"your_endpoint"}, 35 | AccessKey: "your_access_key", 36 | SecretKey: "your_secret_key", 37 | Topic: "your_topic", 38 | Gid: "your_group_id", 39 | Tags: []string{"go-rocketmq"}, 40 | } 41 | 42 | consumer := rocketmq.NewPushConsumer(config) 43 | consumer.Subscribe(consumeMsg) 44 | 45 | consumer.Start() 46 | time.Sleep(time.Hour) 47 | } 48 | 49 | func consumeMsg(ctx context.Context, msg *rocketmq.MessageExt) error { 50 | fmt.Println("msg", msg.String()) 51 | return nil 52 | } 53 | -------------------------------------------------------------------------------- /pkg/broker/rocketmq/examples/producer/producer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package main 20 | 21 | import ( 22 | "context" 23 | "time" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/utils/xstring" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/log" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/broker/rocketmq" 30 | ) 31 | 32 | func main() { 33 | defer log.New(nil).Sync() 34 | 35 | config := &rocketmq.ProducerConfig{ 36 | Endpoint: []string{"your_endpoint"}, 37 | AccessKey: "your_access_key", 38 | SecretKey: "your_secret_key", 39 | Retry: 1, 40 | SendTimeout: time.Second, 41 | Topic: "your_topic", 42 | } 43 | 44 | p := rocketmq.NewProducer(config) 45 | p.Start() 46 | 47 | for i := 0; i < 100; i++ { 48 | p.SendSyncMsg(context.Background(), "Hello RocketMQ Go Client!"+xstring.RandomString(16)) 49 | } 50 | 51 | p.Shutdown() 52 | } 53 | -------------------------------------------------------------------------------- /pkg/broker/rocketmq/primitive.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2022 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package rocketmq 20 | 21 | import ( 22 | "github.com/apache/rocketmq-client-go/v2/primitive" 23 | "github.com/apache/rocketmq-client-go/v2/rlog" 24 | ) 25 | 26 | type ( 27 | // MessageExt is an alias of primitive.MessageExt 28 | MessageExt = primitive.MessageExt 29 | // Message is an alias of primitive.Message 30 | Message = primitive.Message 31 | // Interceptor is an alias of primitive.Interceptor 32 | Interceptor = primitive.Interceptor 33 | ) 34 | 35 | // SetLogLevel set rocket mq log level 36 | func SetLogLevel(level string) { 37 | if level == "" { 38 | return 39 | } 40 | rlog.SetLogLevel(level) 41 | } 42 | -------------------------------------------------------------------------------- /pkg/conf/parser/toml/toml.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package toml 20 | 21 | import "github.com/BurntSushi/toml" 22 | 23 | type TOML map[string]interface{} 24 | 25 | func NewTOMLParser() TOML { 26 | parser := make(TOML) 27 | return parser 28 | } 29 | 30 | // Marshal marshal TOML to bytes 31 | func (t TOML) Marshal(m map[string]interface{}) ([]byte, error) { 32 | return nil, nil 33 | } 34 | 35 | // Unmarshal unmarshal input bytes to map[string]interface{} 36 | func (t TOML) Unmarshal(b []byte) (map[string]interface{}, error) { 37 | if err := toml.Unmarshal(b, &t); err != nil { 38 | return nil, err 39 | } 40 | 41 | return t, nil 42 | } 43 | -------------------------------------------------------------------------------- /pkg/database/mongo/aggregate.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package mongo 20 | 21 | import ( 22 | "context" 23 | "time" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/breaker" 26 | 27 | "github.com/opentracing/opentracing-go" 28 | 29 | "github.com/opentracing/opentracing-go/ext" 30 | "github.com/opentracing/opentracing-go/log" 31 | "github.com/qiniu/qmgo" 32 | ) 33 | 34 | type Aggregate struct { 35 | ai qmgo.AggregateI 36 | ctx context.Context 37 | config *Config 38 | span opentracing.Span 39 | brk *breaker.BreakerGroup 40 | } 41 | 42 | // All iterates the cursor from aggregate and decodes each document into results. 43 | func (a *Aggregate) All(results interface{}) (err error) { 44 | err = a.brk.Do(a.config.Addr, func() error { 45 | now := time.Now() 46 | a.span = a.span.SetOperationName("aggregate_all") 47 | defer a.span.Finish() 48 | 49 | err = a.ai.All(results) 50 | if ok, elapse := slowLog(now, a.config.SlowQueryDuration); ok { 51 | ext.Error.Set(a.span, true) 52 | a.span.LogFields(log.String("event", "slow_query"), log.Int64("elapse", int64(elapse))) 53 | } 54 | return err 55 | }, accept) 56 | return 57 | } 58 | 59 | // One iterates the cursor from aggregate and decodes current document into result. 60 | func (a *Aggregate) One(result interface{}) (err error) { 61 | err = a.brk.Do(a.config.Addr, func() error { 62 | now := time.Now() 63 | a.span = a.span.SetOperationName("aggregate_one") 64 | defer a.span.Finish() 65 | 66 | err = a.ai.One(result) 67 | if ok, elapse := slowLog(now, a.config.SlowQueryDuration); ok { 68 | ext.Error.Set(a.span, true) 69 | a.span.LogFields(log.String("event", "slow_query"), log.Int64("elapse", int64(elapse))) 70 | } 71 | return err 72 | }, accept) 73 | return 74 | } 75 | 76 | // Iter return the cursor after aggregate 77 | // In most scenario do not use Iter 78 | func (a *Aggregate) Iter() qmgo.CursorI { 79 | return a.ai.Iter() 80 | } 81 | -------------------------------------------------------------------------------- /pkg/database/redis/lock_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package redis 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestLock(t *testing.T) { 28 | locked, lockVal, err := r.Lock(ctx, "lock", 10000) 29 | assert.Nil(t, err) 30 | assert.True(t, locked) 31 | locked, _, err = r.Lock(ctx, "lock", 10000) 32 | assert.Nil(t, err) 33 | assert.False(t, locked) 34 | err = r.UnLock(ctx, "lock", lockVal) 35 | assert.Nil(t, err) 36 | err = r.ForceUnLock(ctx, "lock") 37 | assert.Nil(t, err) 38 | } 39 | -------------------------------------------------------------------------------- /pkg/database/sql/mysql.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package sql 20 | 21 | import ( 22 | "github.com/UnderTreeTech/waterdrop/pkg/log" 23 | 24 | _ "github.com/go-sql-driver/mysql" 25 | ) 26 | 27 | // NewMySQL new db and retry connection when has error. 28 | func NewMySQL(c *Config) (db *DB) { 29 | if c.QueryTimeout == 0 || c.ExecTimeout == 0 || c.TranTimeout == 0 { 30 | panic("mysql must be set query/execute/transction timeout") 31 | } 32 | db, err := Open(c) 33 | if err != nil { 34 | log.Errorf("open mysql error", log.Any("err_msg", err)) 35 | panic(err) 36 | } 37 | return 38 | } 39 | -------------------------------------------------------------------------------- /pkg/database/sql/postgres.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package sql 20 | 21 | import ( 22 | "github.com/UnderTreeTech/waterdrop/pkg/log" 23 | _ "github.com/lib/pq" 24 | ) 25 | 26 | // NewPostgres new db and retry connection when has error. 27 | func NewPostgres(c *Config) (db *DB) { 28 | if c.QueryTimeout == 0 || c.ExecTimeout == 0 || c.TranTimeout == 0 { 29 | panic("postgres must be set query/execute/transction timeout") 30 | } 31 | db, err := Open(c) 32 | if err != nil { 33 | log.Errorf("open postgres error", log.Any("err_msg", err)) 34 | panic(err) 35 | } 36 | return 37 | } 38 | -------------------------------------------------------------------------------- /pkg/log/log_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package log 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | func TestLog(t *testing.T) { 29 | type testUser struct { 30 | Password string 31 | Pwd string `json:"passWord"` 32 | AccessToken string `json:"access_token"` 33 | ApiKey string `json:"api_key"` 34 | ApiSecret string `json:"api_secret"` 35 | } 36 | 37 | type testStruct struct { 38 | TestField string `json:"field"` 39 | Hello string `json:"token"` 40 | Users []*testUser `json:"user"` 41 | } 42 | 43 | defaultLogger = newLogger(defaultConfig()) 44 | 45 | Info(context.Background(), "info", 46 | Int64("age", 10), 47 | String("hello", "world"), 48 | Any("any", []string{"shanghai", "xuhui"}), 49 | ) 50 | 51 | Warn(context.Background(), "warn", 52 | String("john", "sun"), 53 | ) 54 | 55 | Debug(context.Background(), "debug", 56 | String("shanghai", "xuhui"), 57 | ) 58 | 59 | Error(context.Background(), "division zero", String("shanghai", "xuhui")) 60 | 61 | assert.Panics(t, func() { 62 | Panic(context.Background(), "memory leaky", String("stop", "yes")) 63 | }) 64 | 65 | Infof("info", 66 | Int64("age", 10), 67 | String("hello", "world"), 68 | Any("any", []string{"shanghai", "xuhui"}), 69 | ) 70 | 71 | Warnf("warn", 72 | String("john", "sun"), 73 | ) 74 | 75 | Debugf("debug", 76 | String("shanghai", "xuhui"), 77 | ) 78 | 79 | Errorf("division zero", String("shanghai", "xuhui")) 80 | 81 | ts := testStruct{ 82 | TestField: "fieldValue", 83 | Hello: "world", 84 | } 85 | ts.Users = append(ts.Users, &testUser{ 86 | Password: "123", 87 | Pwd: "johnsun", 88 | AccessToken: "hello", 89 | ApiKey: "world", 90 | ApiSecret: "yes", 91 | }) 92 | 93 | Infof("filter keyword", 94 | Any("users", ts), 95 | Any("seCret", []string{"shanghai", "xuhui"}), 96 | String("token", "world"), 97 | ) 98 | 99 | assert.Panics(t, func() { 100 | Panicf("memory leaky", String("stop", "yes")) 101 | }) 102 | } 103 | -------------------------------------------------------------------------------- /pkg/ratelimit/setinel/sentinel.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package setinel 20 | 21 | import ( 22 | "encoding/json" 23 | "io/ioutil" 24 | 25 | "github.com/alibaba/sentinel-golang/api" 26 | 27 | sc "github.com/alibaba/sentinel-golang/core/config" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/log" 30 | 31 | "github.com/alibaba/sentinel-golang/core/flow" 32 | ) 33 | 34 | // Config sentinel polices and config 35 | type Config struct { 36 | AppName string 37 | LogPath string 38 | FlowRules []*flow.Rule 39 | RulePath string 40 | } 41 | 42 | // InitSentinel init sentinel by config 43 | func InitSentinel(config *Config) error { 44 | if config.RulePath != "" { 45 | var rules []*flow.Rule 46 | content, err := ioutil.ReadFile(config.RulePath) 47 | if err != nil { 48 | log.Errorf("read rule fail", log.String("rule_path", config.RulePath), log.String("error", err.Error())) 49 | } 50 | 51 | if err := json.Unmarshal(content, &rules); err != nil { 52 | log.Errorf("unmarshal rule fail", log.String("rule_path", config.RulePath), log.String("error", err.Error())) 53 | } 54 | 55 | config.FlowRules = append(config.FlowRules, rules...) 56 | } 57 | 58 | entity := sc.NewDefaultConfig() 59 | entity.Sentinel.App.Name = config.AppName 60 | entity.Sentinel.Log.Dir = config.LogPath 61 | 62 | if len(config.FlowRules) > 0 { 63 | if _, err := flow.LoadRules(config.FlowRules); err != nil { 64 | log.Errorf("load rule fail", log.String("rule_path", config.RulePath), log.String("error", err.Error())) 65 | } 66 | } 67 | 68 | return api.InitWithConfig(entity) 69 | } 70 | -------------------------------------------------------------------------------- /pkg/ratelimit/setinel/sentinel_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package setinel 20 | 21 | import ( 22 | "fmt" 23 | "math/rand" 24 | "testing" 25 | "time" 26 | 27 | "github.com/alibaba/sentinel-golang/api" 28 | 29 | "github.com/alibaba/sentinel-golang/core/base" 30 | "github.com/alibaba/sentinel-golang/util" 31 | 32 | "github.com/alibaba/sentinel-golang/core/flow" 33 | ) 34 | 35 | func TestInitSentinel(t *testing.T) { 36 | config := &Config{ 37 | AppName: "sentinel", 38 | FlowRules: make([]*flow.Rule, 0), 39 | } 40 | 41 | config.FlowRules = append(config.FlowRules, &flow.Rule{ 42 | Resource: "sentinel", 43 | TokenCalculateStrategy: flow.Direct, 44 | ControlBehavior: flow.Reject, 45 | Threshold: 10, 46 | StatIntervalInMs: 1000, 47 | }) 48 | 49 | InitSentinel(config) 50 | 51 | for i := 0; i < 10; i++ { 52 | go func() { 53 | for { 54 | e, b := api.Entry("sentinel", api.WithTrafficType(base.Inbound)) 55 | if b != nil { 56 | // Blocked. We could get the block reason from the BlockError. 57 | fmt.Println("blocked request", b.BlockMsg()) 58 | time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond) 59 | } else { 60 | // Passed, wrap the logic here. 61 | fmt.Println(util.CurrentTimeMillis(), "passed") 62 | time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond) 63 | 64 | // Be sure the entry is exited finally. 65 | e.Exit() 66 | } 67 | } 68 | }() 69 | } 70 | 71 | time.Sleep(time.Second * 1) 72 | } 73 | -------------------------------------------------------------------------------- /pkg/registry/etcd/etcd_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package etcd 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | "time" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/log" 27 | 28 | "github.com/stretchr/testify/assert" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/registry" 31 | ) 32 | 33 | // TestEtcdRegistry test etcd registry 34 | func TestEtcdRegistry(t *testing.T) { 35 | defer log.New(nil).Sync() 36 | etcd := New(defaultConfig) 37 | defer etcd.Close() 38 | service := ®istry.ServiceInfo{ 39 | Name: "service.waterdrop.v1", 40 | Scheme: schemeGRPC, 41 | Addr: "127.0.0.1:9999", 42 | Version: "v1.2", 43 | } 44 | 45 | err := etcd.Register(context.Background(), service) 46 | assert.Nil(t, err) 47 | time.Sleep(time.Millisecond * 10) 48 | services, err := etcd.List(context.Background(), service.Name, "") 49 | assert.Nil(t, err) 50 | assert.Equal(t, len(services), 1) 51 | err = etcd.DeRegister(context.Background(), service) 52 | assert.Nil(t, err) 53 | services, err = etcd.List(context.Background(), service.Name, "") 54 | assert.Nil(t, err) 55 | assert.Equal(t, len(services), 0) 56 | } 57 | -------------------------------------------------------------------------------- /pkg/registry/etcd/lock.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package etcd 20 | 21 | import ( 22 | "context" 23 | "time" 24 | 25 | "go.etcd.io/etcd/client/v3/concurrency" 26 | ) 27 | 28 | // mutex distributed lock based on etcd 29 | type mutex struct { 30 | s *concurrency.Session 31 | m *concurrency.Mutex 32 | } 33 | 34 | // NewMutex new lock 35 | func (er *EtcdRegistry) NewMutex(key string, opts ...concurrency.SessionOption) (m *mutex, err error) { 36 | m = &mutex{} 37 | 38 | // default session ttl 60s 39 | m.s, err = concurrency.NewSession(er.client, opts...) 40 | if err != nil { 41 | return 42 | } 43 | m.m = concurrency.NewMutex(m.s, key) 44 | 45 | return 46 | } 47 | 48 | // Lock do lock op 49 | func (m *mutex) Lock(ctx context.Context, timeout time.Duration) (err error) { 50 | ctx, cancel := context.WithTimeout(ctx, timeout) 51 | defer cancel() 52 | 53 | return m.m.Lock(ctx) 54 | } 55 | 56 | // Unlock release locked resource 57 | func (m *mutex) Unlock(ctx context.Context) (err error) { 58 | err = m.m.Unlock(ctx) 59 | if err != nil { 60 | return 61 | } 62 | 63 | return m.s.Close() 64 | } 65 | -------------------------------------------------------------------------------- /pkg/registry/etcd/lock_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package etcd 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | "time" 25 | 26 | "github.com/stretchr/testify/assert" 27 | ) 28 | 29 | func TestLock(t *testing.T) { 30 | etcd := New(defaultConfig) 31 | defer etcd.Close() 32 | 33 | lock, err := etcd.NewMutex("/waterdrop/test/lock") 34 | assert.Nil(t, err) 35 | err = lock.Lock(context.Background(), time.Second) 36 | assert.Nil(t, err) 37 | defer lock.Unlock(context.Background()) 38 | 39 | lock2, err := etcd.NewMutex("/waterdrop/test/lock") 40 | defer lock2.Unlock(context.Background()) 41 | assert.Nil(t, err) 42 | err = lock2.Lock(context.Background(), time.Second) 43 | assert.NotNil(t, err) 44 | } 45 | -------------------------------------------------------------------------------- /pkg/registry/registry.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package registry 20 | 21 | import "context" 22 | 23 | // metadata common key 24 | const ( 25 | MetaWeight = "weight" 26 | MetaCluster = "cluster" 27 | MetaZone = "zone" 28 | MetaColor = "color" 29 | ) 30 | 31 | type Registry interface { 32 | Register(ctx context.Context, info *ServiceInfo) error 33 | DeRegister(ctx context.Context, info *ServiceInfo) error 34 | List(ctx context.Context, name string, scheme string) (services []*ServiceInfo, err error) 35 | Close() 36 | } 37 | 38 | // ServiceInfo service metadata definition 39 | type ServiceInfo struct { 40 | // Service Name 41 | Name string `json:"name"` 42 | // Service Scheme, http/grpc 43 | Scheme string `json:"schema"` 44 | // Service Addr 45 | Addr string `json:"addr"` 46 | // Metadata is the information associated with Addr, which may be used 47 | // to make load balancing decision 48 | Metadata map[string]string `json:"metadata"` 49 | // Region is region 50 | Region string `json:"region"` 51 | // Zone is IDC 52 | Zone string `json:"zone"` 53 | // prod/pre/test/dev 54 | Env string `json:"env"` 55 | // Service Version 56 | Version string `json:"version"` 57 | } 58 | -------------------------------------------------------------------------------- /pkg/server/http/config/config.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package config 20 | 21 | import ( 22 | "time" 23 | 24 | "github.com/gin-gonic/gin" 25 | ) 26 | 27 | // ClientConfig http client config 28 | type ClientConfig struct { 29 | // HostURL peer service host 30 | HostURL string 31 | // Timeout request timeout 32 | Timeout time.Duration 33 | // SlowRequestDuration slow request timeout 34 | SlowRequestDuration time.Duration 35 | // EnableDebug trace request details 36 | EnableDebug bool 37 | // Key client key 38 | Key string 39 | // Secret signature secret 40 | Secret string 41 | } 42 | 43 | // ServerConfig http server config 44 | type ServerConfig struct { 45 | // Addr server addr, like :8080 or 127.0.0.1:8080 46 | Addr string 47 | // Timeout request timeout 48 | Timeout time.Duration 49 | // Mode server mode: release or debug 50 | Mode string 51 | // SlowRequestDuration slow request timeout 52 | SlowRequestDuration time.Duration 53 | // WatchConfig whether watch config file changes 54 | WatchConfig bool 55 | } 56 | 57 | // DefaultServerConfig default server configs, for start http server out of box 58 | func DefaultServerConfig() *ServerConfig { 59 | return &ServerConfig{ 60 | Addr: "0.0.0.0:10000", 61 | Mode: gin.ReleaseMode, 62 | Timeout: time.Millisecond * 1000, 63 | SlowRequestDuration: 500 * time.Millisecond, 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /pkg/server/http/metadata/metadata.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metadata 20 | 21 | import ( 22 | "net/http" 23 | "strconv" 24 | "time" 25 | ) 26 | 27 | const ( 28 | HeaderContentType = "Content-Type" 29 | HeaderUserAgent = "User-Agent" 30 | HeaderAppkey = "Appkey" 31 | HeaderTimestamp = "Timestamp" 32 | HeaderSign = "Sign" 33 | HeaderNonce = "Nonce" 34 | HeaderAcceptLanguage = "Accept-Language" 35 | HeaderHttpTimeout = "X-Request-Timeout" 36 | HeaderHttpTraceId = "X-Trace-Id" 37 | 38 | DefaultContentTypeJson = "application/json;charset=utf-8" 39 | DefaultUserAgentVal = "waterdrop" 40 | DefaultLocale = "zh-CN" 41 | DefaultNonceLen = 16 42 | 43 | LimitBodyBytes = 4096 44 | ) 45 | 46 | // GetTimeout get timeout from request header 47 | // similar as grpc 48 | func GetTimeout(req *http.Request) time.Duration { 49 | to := req.Header.Get(HeaderHttpTimeout) 50 | timeout, err := strconv.ParseInt(to, 10, 64) 51 | //reduce 5ms network transmission time for every request 52 | if err == nil && timeout > 5 { 53 | timeout -= 5 54 | } 55 | 56 | return time.Duration(timeout) * time.Millisecond 57 | } 58 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/cors.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "github.com/gin-contrib/cors" 23 | "github.com/gin-gonic/gin" 24 | ) 25 | 26 | // CORSConfig is an alias of cors.Config 27 | type CORSConfig cors.Config 28 | 29 | // DefaultCORS default cors handler 30 | func DefaultCORS() gin.HandlerFunc { 31 | return cors.Default() 32 | } 33 | 34 | // NewCORS customer cors handler by config 35 | func NewCORS(config CORSConfig) gin.HandlerFunc { 36 | return cors.New(cors.Config(config)) 37 | } 38 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/logger.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "encoding/json" 23 | "time" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/log" 28 | 29 | "github.com/gin-gonic/gin" 30 | ) 31 | 32 | // Logger log request details 33 | func Logger(config *config.ServerConfig) gin.HandlerFunc { 34 | return func(c *gin.Context) { 35 | now := time.Now() 36 | var quota float64 37 | if deadline, ok := c.Request.Context().Deadline(); ok { 38 | quota = time.Until(deadline).Seconds() 39 | } 40 | 41 | c.Next() 42 | 43 | duration := time.Since(now) 44 | 45 | fields := make([]log.Field, 0, 10) 46 | fields = append( 47 | fields, 48 | log.String("client_ip", c.ClientIP()), 49 | log.String("method", c.Request.Method), 50 | log.String("path", c.Request.URL.Path), 51 | log.Any("headers", c.Request.Header), 52 | log.Any("req", json.RawMessage(log.JsonForm(c.Request.Form))), 53 | log.Float64("quota", quota), 54 | log.Float64("duration", duration.Seconds()), 55 | log.Int("status", c.Writer.Status()), 56 | log.Int("size", c.Writer.Size()), 57 | log.String("error", c.Errors.ByType(gin.ErrorTypePrivate).String()), 58 | ) 59 | 60 | if duration >= config.SlowRequestDuration { 61 | log.Warn(c.Request.Context(), "http-slow-access-log", fields...) 62 | } else { 63 | log.Info(c.Request.Context(), "http-access-log", fields...) 64 | } 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/logger_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "net/http" 23 | "net/http/httptest" 24 | "testing" 25 | "time" 26 | 27 | "github.com/stretchr/testify/assert" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 30 | "github.com/gin-gonic/gin" 31 | ) 32 | 33 | func TestLogger(t *testing.T) { 34 | engine := gin.New() 35 | engine.Use(Logger(config.DefaultServerConfig())) 36 | 37 | engine.GET("/log/normal", func(ctx *gin.Context) { 38 | ctx.String(http.StatusOK, "normal") 39 | }) 40 | engine.GET("/log/slow", func(ctx *gin.Context) { 41 | time.Sleep(time.Second) 42 | ctx.String(http.StatusOK, "slow") 43 | }) 44 | 45 | var tests = []string{"/log/normal", "/log/slow"} 46 | 47 | for _, test := range tests { 48 | req := httptest.NewRequest(http.MethodGet, test, nil) 49 | w := httptest.NewRecorder() 50 | engine.ServeHTTP(w, req) 51 | 52 | assert.Contains(t, test, w.Body.String()) 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/metric.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "strconv" 23 | "time" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/stats/metric" 26 | "github.com/gin-gonic/gin" 27 | ) 28 | 29 | var namespace = "appkey" 30 | 31 | const _defaultNamespace = "default" 32 | 33 | // Metric http request metric middleware 34 | func Metric() gin.HandlerFunc { 35 | return func(c *gin.Context) { 36 | now := time.Now() 37 | 38 | c.Next() 39 | 40 | ns := c.Request.Header.Get(namespace) 41 | if ns == "" { 42 | ns = _defaultNamespace 43 | } 44 | metric.HTTPServerHandleCounter.Inc(c.FullPath(), c.Request.Method, ns, strconv.Itoa(c.Writer.Status())) 45 | metric.HTTPServerReqDuration.Observe(time.Since(now).Seconds(), c.FullPath(), c.Request.Method, ns) 46 | } 47 | } 48 | 49 | // SetHttpMetricNamespace set http metric namespace get from which http header 50 | // default namespace is get from http header: appkey 51 | func SetMetricNamespace(header string) { 52 | namespace = header 53 | } 54 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/metric_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "net/http" 23 | "net/http/httptest" 24 | "testing" 25 | 26 | "github.com/gin-gonic/gin" 27 | "github.com/stretchr/testify/assert" 28 | ) 29 | 30 | func TestMetric(t *testing.T) { 31 | engine := gin.New() 32 | engine.Use(Metric()) 33 | engine.GET("/", func(ctx *gin.Context) { 34 | ctx.String(http.StatusOK, "normal") 35 | }) 36 | 37 | req := httptest.NewRequest(http.MethodGet, "/", nil) 38 | w := httptest.NewRecorder() 39 | engine.ServeHTTP(w, req) 40 | 41 | assert.Equal(t, http.StatusOK, w.Code) 42 | } 43 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/ratelimit/limiter.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package ratelimit 20 | 21 | import ( 22 | "context" 23 | 24 | "github.com/gin-gonic/gin" 25 | ) 26 | 27 | type Limiter interface { 28 | Allow(context.Context, string) bool 29 | } 30 | 31 | type Option func(*LimitOpts) 32 | 33 | type LimitOpts struct { 34 | Strategy func(c *gin.Context) string 35 | Fallback func(c *gin.Context) 36 | } 37 | 38 | // Apply applies the option to limiter 39 | func Apply(opts []Option) *LimitOpts { 40 | limitOpts := &LimitOpts{} 41 | for _, opt := range opts { 42 | opt(limitOpts) 43 | } 44 | 45 | return limitOpts 46 | } 47 | 48 | // WithFallback set the fallback handler when request is blocked 49 | func WithFallback(fallback func(c *gin.Context)) Option { 50 | return func(opts *LimitOpts) { 51 | opts.Fallback = fallback 52 | } 53 | } 54 | 55 | // WithResourceStrategy set the verify path extract strategy of request 56 | func WithResourceStrategy(strategy func(c *gin.Context) string) Option { 57 | return func(opts *LimitOpts) { 58 | opts.Strategy = strategy 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/ratelimit/sentinel/sentinel.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package sentinel 20 | 21 | import ( 22 | "net/http" 23 | 24 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/middlewares/ratelimit" 25 | 26 | "github.com/alibaba/sentinel-golang/api" 27 | "github.com/alibaba/sentinel-golang/core/base" 28 | "github.com/gin-gonic/gin" 29 | ) 30 | 31 | // Sentinel return rate limit middleware 32 | func Sentinel(opts ...ratelimit.Option) gin.HandlerFunc { 33 | limitOption := ratelimit.Apply(opts) 34 | return func(c *gin.Context) { 35 | limitPath := c.Request.Method + ":" + c.FullPath() 36 | if limitOption.Strategy != nil { 37 | limitPath = limitOption.Strategy(c) 38 | } 39 | 40 | entry, err := api.Entry(limitPath, api.WithResourceType(base.ResTypeWeb), api.WithTrafficType(base.Inbound)) 41 | if err != nil { 42 | if limitOption.Fallback != nil { 43 | limitOption.Fallback(c) 44 | } else { 45 | c.AbortWithStatus(http.StatusTooManyRequests) 46 | } 47 | return 48 | } 49 | 50 | defer entry.Exit() 51 | c.Next() 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/recovery_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "net" 23 | "net/http" 24 | "net/http/httptest" 25 | "os" 26 | "syscall" 27 | "testing" 28 | 29 | "github.com/stretchr/testify/assert" 30 | 31 | "github.com/gin-gonic/gin" 32 | 33 | "github.com/UnderTreeTech/waterdrop/pkg/log" 34 | ) 35 | 36 | func TestMain(m *testing.M) { 37 | defer log.New(nil).Sync() 38 | 39 | code := m.Run() 40 | os.Exit(code) 41 | } 42 | 43 | func TestRecovery(t *testing.T) { 44 | engine := gin.New() 45 | engine.Use(Recovery()) 46 | 47 | engine.GET("/recovery/panic", func(ctx *gin.Context) { 48 | panic("internal server error occurred") 49 | }) 50 | 51 | req := httptest.NewRequest(http.MethodGet, "/recovery/panic", nil) 52 | w := httptest.NewRecorder() 53 | engine.ServeHTTP(w, req) 54 | 55 | assert.Equal(t, http.StatusInternalServerError, w.Code) 56 | } 57 | 58 | func TestPanicWithAbort(t *testing.T) { 59 | engine := gin.New() 60 | engine.Use(Recovery()) 61 | 62 | engine.GET("/recovery/panic", func(ctx *gin.Context) { 63 | ctx.AbortWithStatus(http.StatusBadRequest) 64 | panic("internal server error occurred") 65 | }) 66 | 67 | req := httptest.NewRequest(http.MethodGet, "/recovery/panic", nil) 68 | w := httptest.NewRecorder() 69 | engine.ServeHTTP(w, req) 70 | 71 | assert.Equal(t, http.StatusBadRequest, w.Code) 72 | } 73 | 74 | func TestPanicWithBrokenPipe(t *testing.T) { 75 | const expectCode = 204 76 | 77 | expectMsgs := map[syscall.Errno]string{ 78 | syscall.EPIPE: "broken pipe", 79 | syscall.ECONNRESET: "connection reset by peer", 80 | } 81 | 82 | for errno, expectMsg := range expectMsgs { 83 | t.Run(expectMsg, func(t *testing.T) { 84 | engine := gin.New() 85 | engine.Use(Recovery()) 86 | engine.GET("/recovery/panic", func(ctx *gin.Context) { 87 | // Start writing response 88 | ctx.Header("X-Test", "Value") 89 | ctx.Status(expectCode) 90 | 91 | // Client connection closed 92 | e := &net.OpError{Err: &os.SyscallError{Err: errno}} 93 | panic(e) 94 | }) 95 | // RUN 96 | req := httptest.NewRequest(http.MethodGet, "/recovery/panic", nil) 97 | w := httptest.NewRecorder() 98 | engine.ServeHTTP(w, req) 99 | assert.Equal(t, expectCode, w.Code) 100 | }) 101 | } 102 | } 103 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/trace.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "context" 23 | 24 | "github.com/UnderTreeTech/waterdrop/pkg/trace" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/metadata" 27 | "github.com/opentracing/opentracing-go/ext" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 30 | 31 | "github.com/gin-gonic/gin" 32 | ) 33 | 34 | // Trace trace incoming request details 35 | func Trace(config *config.ServerConfig) gin.HandlerFunc { 36 | return func(c *gin.Context) { 37 | span, ctx := trace.StartSpanFromContext( 38 | c.Request.Context(), 39 | c.Request.Method+" "+c.Request.URL.Path, 40 | trace.HeaderExtractor(c.Request.Header), 41 | ) 42 | ext.Component.Set(span, "http") 43 | ext.SpanKind.Set(span, ext.SpanKindRPCServerEnum) 44 | ext.HTTPMethod.Set(span, c.Request.Method) 45 | ext.HTTPUrl.Set(span, c.FullPath()) 46 | ext.PeerHostIPv4.SetString(span, c.ClientIP()) 47 | 48 | // adjust request timeout 49 | timeout := config.Timeout 50 | reqTimeout := metadata.GetTimeout(c.Request) 51 | if reqTimeout > 0 && timeout > reqTimeout { 52 | timeout = reqTimeout 53 | } 54 | 55 | // if zero timeout config means never timeout 56 | var cancel func() 57 | if timeout > 0 { 58 | ctx, cancel = context.WithTimeout(ctx, timeout) 59 | } else { 60 | cancel = func() {} 61 | } 62 | defer func() { 63 | span.Finish() 64 | cancel() 65 | }() 66 | 67 | c.Request = c.Request.WithContext(ctx) 68 | c.Writer.Header().Set(metadata.HeaderHttpTraceId, trace.TraceID(ctx)) 69 | 70 | c.Next() 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /pkg/server/http/middlewares/trace_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package middlewares 20 | 21 | import ( 22 | "fmt" 23 | "net/http" 24 | "net/http/httptest" 25 | "testing" 26 | 27 | "github.com/stretchr/testify/assert" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/trace" 30 | 31 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 32 | "github.com/gin-gonic/gin" 33 | 34 | opentracing "github.com/opentracing/opentracing-go" 35 | jconfig "github.com/uber/jaeger-client-go/config" 36 | ) 37 | 38 | func newJaegerClient() (opentracing.Tracer, func()) { 39 | var configuration = jconfig.Configuration{ 40 | ServiceName: "trace", 41 | } 42 | 43 | tracer, closer, err := configuration.NewTracer() 44 | if err != nil { 45 | panic(fmt.Sprintf("new jaeger trace fail, err msg %s", err.Error())) 46 | } 47 | 48 | return tracer, func() { closer.Close() } 49 | } 50 | 51 | func TestTrace(t *testing.T) { 52 | engine := gin.New() 53 | 54 | engine.Use(Trace(config.DefaultServerConfig())) 55 | engine.GET("/trace/mock", func(ctx *gin.Context) { 56 | ctx.String(http.StatusOK, trace.TraceID(ctx.Request.Context())) 57 | }) 58 | 59 | req := httptest.NewRequest(http.MethodGet, "/trace/mock", nil) 60 | w := httptest.NewRecorder() 61 | engine.ServeHTTP(w, req) 62 | 63 | assert.Equal(t, w.Body.String(), "") 64 | } 65 | 66 | func TestJaegerTrace(t *testing.T) { 67 | engine := gin.New() 68 | engine.Use(Trace(config.DefaultServerConfig())) 69 | engine.GET("/trace/jaeger", func(ctx *gin.Context) { 70 | ctx.String(http.StatusOK, trace.TraceID(ctx.Request.Context())) 71 | }) 72 | 73 | tracer, close := newJaegerClient() 74 | trace.SetGlobalTracer(tracer) 75 | defer close() 76 | 77 | req := httptest.NewRequest(http.MethodGet, "/trace/jaeger", nil) 78 | w := httptest.NewRecorder() 79 | engine.ServeHTTP(w, req) 80 | 81 | assert.NotEqual(t, 0, len(w.Body.String())) 82 | } 83 | -------------------------------------------------------------------------------- /pkg/server/http/server/server.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package server 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "log" 25 | "net" 26 | "net/http" 27 | 28 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/websocket" 31 | 32 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/middlewares" 33 | 34 | "github.com/gin-gonic/gin" 35 | 36 | _ "github.com/UnderTreeTech/waterdrop/pkg/version" 37 | // Automatically set GOMAXPROCS to match Linux container CPU quota 38 | _ "go.uber.org/automaxprocs" 39 | ) 40 | 41 | // Server http server 42 | type Server struct { 43 | *gin.Engine 44 | Server *http.Server 45 | config *config.ServerConfig 46 | } 47 | 48 | // New returns a http server instance 49 | func New(cfg *config.ServerConfig) *Server { 50 | if cfg == nil { 51 | cfg = config.DefaultServerConfig() 52 | } 53 | 54 | gin.SetMode(cfg.Mode) 55 | srv := &Server{ 56 | Engine: gin.New(), 57 | config: cfg, 58 | } 59 | 60 | srv.Use(middlewares.Recovery(), middlewares.Trace(srv.config), middlewares.Logger(srv.config), middlewares.Metric()) 61 | return srv 62 | } 63 | 64 | // Start server 65 | func (s *Server) Start() net.Addr { 66 | listener, err := net.Listen("tcp", s.config.Addr) 67 | if err != nil { 68 | panic(fmt.Sprintf("http server: listen tcp fail,err msg %s", err.Error())) 69 | } 70 | 71 | s.Server = &http.Server{ 72 | Addr: s.config.Addr, 73 | Handler: s, 74 | } 75 | 76 | go func() { 77 | if err := s.Server.Serve(listener); err != nil { 78 | if err == http.ErrServerClosed { 79 | log.Printf("waterdrop: http server closed") 80 | return 81 | } 82 | panic(fmt.Sprintf("HTTP Server serve fail,err msg %s", err.Error())) 83 | } 84 | }() 85 | 86 | log.Printf("HTTP Server: start http listen addr: %s", listener.Addr().String()) 87 | return listener.Addr() 88 | } 89 | 90 | // Stop shutdown server graceful 91 | func (s *Server) Stop(ctx context.Context) error { 92 | return s.Server.Shutdown(ctx) 93 | } 94 | 95 | // Upgrade upgrade http to websocket 96 | func (s *Server) Upgrade(ws *websocket.WebSocket) gin.IRoutes { 97 | return s.GET(ws.Path, func(c *gin.Context) { 98 | ws.Upgrade(c.Writer, c.Request) 99 | }) 100 | } 101 | -------------------------------------------------------------------------------- /pkg/server/http/server/server_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package server 20 | 21 | import ( 22 | "context" 23 | "os" 24 | "testing" 25 | "time" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/log" 28 | "github.com/stretchr/testify/assert" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/server/http/config" 31 | ) 32 | 33 | var srv = New(config.DefaultServerConfig()) 34 | 35 | func TestMain(m *testing.M) { 36 | defer log.New(nil).Sync() 37 | 38 | code := m.Run() 39 | os.Exit(code) 40 | } 41 | 42 | func TestStart(t *testing.T) { 43 | net := srv.Start() 44 | assert.Equal(t, "[::]:10000", net.String()) 45 | assert.Equal(t, "tcp", net.Network()) 46 | } 47 | 48 | func TestStop(t *testing.T) { 49 | go func() { 50 | time.Sleep(100 * time.Millisecond) 51 | err := srv.Stop(context.Background()) 52 | assert.Equal(t, err, nil) 53 | }() 54 | 55 | time.Sleep(200 * time.Millisecond) 56 | } 57 | -------------------------------------------------------------------------------- /pkg/server/http/websocket/websocket.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package websocket 20 | 21 | import ( 22 | "net/http" 23 | 24 | "github.com/UnderTreeTech/waterdrop/pkg/log" 25 | "github.com/gorilla/websocket" 26 | ) 27 | 28 | // WebSocketHandler ws callback handler 29 | type WebSocketHandler func(*WebSocket) 30 | 31 | // WebSocket ws definition 32 | type WebSocket struct { 33 | Path string 34 | Handler WebSocketHandler 35 | 36 | *websocket.Upgrader 37 | *websocket.Conn 38 | } 39 | 40 | // NewWebSocket returns a WebSocket instance 41 | func NewWebSocket(path string, handler WebSocketHandler) *WebSocket { 42 | return &WebSocket{ 43 | Path: path, 44 | Upgrader: &websocket.Upgrader{}, 45 | Handler: handler, 46 | } 47 | } 48 | 49 | // Upgrade upgrade http to WebSocket 50 | func (ws *WebSocket) Upgrade(w http.ResponseWriter, r *http.Request) { 51 | conn, err := ws.Upgrader.Upgrade(w, r, nil) 52 | if err != nil { 53 | log.Error(r.Context(), "upgrade fail", log.String("error", err.Error())) 54 | return 55 | } 56 | defer conn.Close() 57 | 58 | ws.Conn = conn 59 | ws.Handler(ws) 60 | } 61 | -------------------------------------------------------------------------------- /pkg/server/rpc/client/client_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package client 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | "time" 25 | 26 | "github.com/UnderTreeTech/waterdrop/tests/proto/demo" 27 | "google.golang.org/protobuf/types/known/emptypb" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/log" 30 | 31 | "github.com/stretchr/testify/assert" 32 | 33 | "github.com/UnderTreeTech/waterdrop/pkg/server/rpc/config" 34 | "github.com/UnderTreeTech/waterdrop/pkg/server/rpc/server" 35 | ) 36 | 37 | // TestClient test grpc client 38 | func TestClient(t *testing.T) { 39 | defer log.New(nil).Sync() 40 | srv := server.New(&config.ServerConfig{ 41 | Addr: "0.0.0.0:21819", 42 | }) 43 | demo.RegisterDemoServer(srv.Server(), &service{}) 44 | srv.Start() 45 | time.Sleep(time.Millisecond * 100) 46 | defer srv.Stop(context.Background()) 47 | 48 | cfg := &config.ClientConfig{ 49 | DialTimeout: 150 * time.Millisecond, 50 | Block: false, 51 | Balancer: "round_robin", 52 | Target: "127.0.0.1:21819", 53 | } 54 | client := New(cfg) 55 | rpc := demo.NewDemoClient(client.GetConn()) 56 | reply, err := rpc.SayHelloURL(context.Background(), &demo.HelloReq{Name: "waterdrop"}) 57 | assert.Equal(t, reply.Content, "Hello waterdrop") 58 | assert.Nil(t, err) 59 | } 60 | 61 | // TestDialTimeout test dial timeout 62 | func TestDialTimeout(t *testing.T) { 63 | defer log.New(nil).Sync() 64 | flag := false 65 | defer func() { 66 | if r := recover(); r != nil { 67 | flag = true 68 | } 69 | assert.Equal(t, flag, true) 70 | }() 71 | 72 | srv := server.New(&config.ServerConfig{ 73 | Addr: "0.0.0.0:21819", 74 | }) 75 | demo.RegisterDemoServer(srv.Server(), &service{}) 76 | go func() { 77 | time.Sleep(time.Millisecond * 100) 78 | srv.Start() 79 | }() 80 | defer srv.Stop(context.Background()) 81 | 82 | cfg := &config.ClientConfig{ 83 | DialTimeout: 50 * time.Millisecond, 84 | Block: true, 85 | Balancer: "round_robin", 86 | Target: "127.0.0.1:21819", 87 | } 88 | client := New(cfg) 89 | demo.NewDemoClient(client.GetConn()) 90 | } 91 | 92 | type service struct{} 93 | 94 | func (s *service) SayHello(ctx context.Context, req *demo.HelloReq) (reply *emptypb.Empty, err error) { 95 | reply = &emptypb.Empty{} 96 | return reply, nil 97 | } 98 | 99 | func (s *service) SayHelloURL(ctx context.Context, req *demo.HelloReq) (reply *demo.HelloResp, err error) { 100 | reply = &demo.HelloResp{Content: "Hello " + req.Name} 101 | return 102 | } 103 | -------------------------------------------------------------------------------- /pkg/server/rpc/config/config.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package config 20 | 21 | import ( 22 | "time" 23 | ) 24 | 25 | // ServerConfig rpc server config 26 | type ServerConfig struct { 27 | // Addr server addr,it may be ":8080" or "127.0.0.1:8080" 28 | Addr string 29 | // Timeout rpc request timeout 30 | Timeout time.Duration 31 | // GRPC ServerParameters 32 | IdleTimeout time.Duration 33 | MaxLifeTime time.Duration 34 | ForceCloseWait time.Duration 35 | KeepAliveInterval time.Duration 36 | KeepAliveTimeout time.Duration 37 | // SlowRequestDuration slow rpc request timeout 38 | SlowRequestDuration time.Duration 39 | // WatchConfig whether watch config file changes 40 | WatchConfig bool 41 | // NotLog escape log detail path 42 | NotLog []string 43 | MaxReceiveMessageSize int 44 | } 45 | 46 | // DefaultServerConfig default server config for starting rpc server out of box 47 | func DefaultServerConfig() *ServerConfig { 48 | return &ServerConfig{ 49 | Addr: "0.0.0.0:20812", 50 | Timeout: time.Second, 51 | IdleTimeout: 180 * time.Second, 52 | MaxLifeTime: 2 * time.Hour, 53 | ForceCloseWait: 20 * time.Second, 54 | KeepAliveInterval: 60 * time.Second, 55 | KeepAliveTimeout: 20 * time.Second, 56 | SlowRequestDuration: 500 * time.Millisecond, 57 | } 58 | } 59 | 60 | // ClientConfig rpc client configs 61 | type ClientConfig struct { 62 | // DialTimeout dial rpc server timeout 63 | DialTimeout time.Duration 64 | // Block dial mode: sync or async 65 | Block bool 66 | // Balancer client balancer, default round robbin 67 | Balancer string 68 | // Target rpc server endpoint 69 | Target string 70 | // Timeout rpc request timeout 71 | Timeout time.Duration 72 | // GRPC ClientParameters 73 | KeepAliveInterval time.Duration 74 | KeepAliveTimeout time.Duration 75 | // SlowRequestDuration client slow request timeout 76 | SlowRequestDuration time.Duration 77 | // NotLog escape log detail path 78 | NotLog []string 79 | // MaxCallSendMsgSize default 4*1024*1024 80 | MaxCallSendMsgSize int 81 | } 82 | 83 | // DefaultClientConfig default client config for starting rpc client out of box 84 | func DefaultClientConfig() *ClientConfig { 85 | return &ClientConfig{ 86 | DialTimeout: 5 * time.Second, 87 | Block: true, 88 | Balancer: "round_robin", 89 | 90 | Timeout: 500 * time.Millisecond, 91 | 92 | KeepAliveInterval: 60 * time.Second, 93 | KeepAliveTimeout: 20 * time.Second, 94 | } 95 | } 96 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/breaker.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | 24 | "github.com/UnderTreeTech/waterdrop/pkg/breaker" 25 | "github.com/UnderTreeTech/waterdrop/pkg/status" 26 | 27 | "google.golang.org/grpc" 28 | ) 29 | 30 | // GoogleSREBreaker breaker based on google sre. It rejects request adaptively based on server response 31 | func GoogleSREBreaker(breakers *breaker.BreakerGroup) grpc.UnaryClientInterceptor { 32 | return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { 33 | return breakers.Do( 34 | method, 35 | func() error { 36 | return invoker(ctx, method, req, reply, cc, opts...) 37 | }, 38 | func(err error) bool { 39 | switch status.ExtractStatus(err).Code() { 40 | case status.Deadline.Code(), status.LimitExceed.Code(), 41 | status.ServerErr.Code(), status.Canceled.Code(), 42 | status.ServiceUnavailable.Code(): 43 | return false 44 | default: 45 | return true 46 | } 47 | }) 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/breaker_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/breaker" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/status" 28 | "github.com/stretchr/testify/assert" 29 | 30 | "google.golang.org/grpc" 31 | ) 32 | 33 | func TestGoogleSREBreaker(t *testing.T) { 34 | interceptor := GoogleSREBreaker(breaker.NewBreakerGroup()) 35 | invoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error { 36 | return nil 37 | } 38 | 39 | t.Run("success", func(t *testing.T) { 40 | err := interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 41 | assert.Nil(t, err) 42 | }) 43 | 44 | invoker = func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error { 45 | return status.ServiceUnavailable 46 | } 47 | t.Run("blocked", func(t *testing.T) { 48 | var err error 49 | for i := 0; i < 10; i++ { 50 | err = interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 51 | } 52 | assert.IsType(t, &status.Status{}, err) 53 | }) 54 | } 55 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/logger_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/server/rpc/config" 26 | 27 | "github.com/stretchr/testify/assert" 28 | 29 | "google.golang.org/grpc" 30 | ) 31 | 32 | func TestLoggerForUnaryServer(t *testing.T) { 33 | interceptor := LoggerForUnaryServer(config.DefaultServerConfig()) 34 | handler := func(ctx context.Context, req interface{}) (resp interface{}, err error) { 35 | return 36 | } 37 | info := &grpc.UnaryServerInfo{ 38 | FullMethod: "/grpc.testing.TestService/UnaryCall", 39 | } 40 | 41 | t.Run("server log", func(t *testing.T) { 42 | resp, err := interceptor(context.Background(), nil, info, handler) 43 | assert.Nil(t, resp, nil) 44 | assert.Nil(t, err, nil) 45 | }) 46 | } 47 | 48 | func TestLoggerForUnaryClient(t *testing.T) { 49 | interceptor := LoggerForUnaryClient(config.DefaultClientConfig()) 50 | invoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) (err error) { 51 | return nil 52 | } 53 | 54 | t.Run("client log", func(t *testing.T) { 55 | err := interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 56 | assert.Nil(t, err, nil) 57 | }) 58 | } 59 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/metric.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "time" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/status" 26 | 27 | "google.golang.org/grpc/peer" 28 | 29 | "github.com/UnderTreeTech/waterdrop/pkg/stats/metric" 30 | 31 | "google.golang.org/grpc" 32 | ) 33 | 34 | // Metric metric handler, currently only metric at server side 35 | func Metric() grpc.UnaryServerInterceptor { 36 | return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) { 37 | now := time.Now() 38 | var ip string 39 | if peer, ok := peer.FromContext(ctx); ok { 40 | ip = peer.Addr.String() 41 | } 42 | 43 | // call server interceptor 44 | resp, err = handler(ctx, req) 45 | 46 | estatus := status.ExtractStatus(err) 47 | metric.UnaryServerHandleCounter.Inc(ip, info.FullMethod, estatus.Error()) 48 | metric.UnaryServerReqDuration.Observe(time.Since(now).Seconds(), ip, info.FullMethod) 49 | return 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/metric_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import "testing" 22 | 23 | func TestMetric(t *testing.T) { 24 | 25 | } 26 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/recovery_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "testing" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/server/rpc/config" 27 | 28 | "github.com/UnderTreeTech/waterdrop/pkg/status" 29 | 30 | "github.com/stretchr/testify/assert" 31 | 32 | "google.golang.org/grpc" 33 | ) 34 | 35 | func TestRecoveryForUnaryServer(t *testing.T) { 36 | interceptor := RecoveryForUnaryServer(config.DefaultServerConfig()) 37 | handler := func(ctx context.Context, req interface{}) (resp interface{}, err error) { 38 | mockSlice := []int{1, 2, 3} 39 | fmt.Println(mockSlice[4]) 40 | return 41 | } 42 | info := &grpc.UnaryServerInfo{ 43 | FullMethod: "/grpc.testing.TestService/UnaryCall", 44 | } 45 | 46 | t.Run("recovery", func(t *testing.T) { 47 | resp, err := interceptor(context.Background(), nil, info, handler) 48 | assert.Nil(t, resp) 49 | errStatus := status.ExtractStatus(err) 50 | assert.Equal(t, status.ServerErr.Code(), errStatus.Code()) 51 | }) 52 | } 53 | 54 | func TestRecoveryForUnaryClient(t *testing.T) { 55 | interceptor := RecoveryForUnaryClient(config.DefaultClientConfig()) 56 | invoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) (err error) { 57 | mockSlice := []int{1, 2, 3} 58 | fmt.Println(mockSlice[4]) 59 | return 60 | } 61 | 62 | t.Run("recovery", func(t *testing.T) { 63 | err := interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 64 | errStatus := status.ExtractStatus(err) 65 | assert.Equal(t, status.ServerErr.Code(), errStatus.Code()) 66 | }) 67 | } 68 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/sentinel.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/log" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/status" 28 | 29 | "github.com/alibaba/sentinel-golang/core/base" 30 | 31 | "github.com/alibaba/sentinel-golang/api" 32 | 33 | "google.golang.org/grpc" 34 | ) 35 | 36 | // SentinelForUnaryClient is client side sentinel 37 | func SentinelForUnaryClient() grpc.UnaryClientInterceptor { 38 | return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) { 39 | entry, blockErr := api.Entry(method, api.WithResourceType(base.ResTypeRPC), api.WithTrafficType(base.Outbound)) 40 | if blockErr != nil { 41 | log.Warn(ctx, 42 | "rpc hit rate limit", 43 | log.String("kind", "client"), 44 | log.String("method", method), 45 | log.String("error", blockErr.Error()), 46 | ) 47 | return status.LimitExceed 48 | } 49 | defer entry.Exit() 50 | fmt.Println("method", method) 51 | err = invoker(ctx, method, req, reply, cc, opts...) 52 | return 53 | } 54 | } 55 | 56 | // SentinelForUnaryServer is server side sentinel 57 | func SentinelForUnaryServer() grpc.UnaryServerInterceptor { 58 | return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) { 59 | entry, blockErr := api.Entry(info.FullMethod, api.WithResourceType(base.ResTypeRPC), api.WithTrafficType(base.Inbound)) 60 | if blockErr != nil { 61 | log.Warn(ctx, 62 | "rpc hit rate limit", 63 | log.String("kind", "server"), 64 | log.String("method", info.FullMethod), 65 | log.String("error", blockErr.Error()), 66 | ) 67 | return nil, status.LimitExceed 68 | } 69 | defer entry.Exit() 70 | 71 | resp, err = handler(ctx, req) 72 | return 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/trace.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | 24 | "google.golang.org/grpc/metadata" 25 | 26 | "google.golang.org/grpc/peer" 27 | 28 | "github.com/opentracing/opentracing-go/ext" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/status" 31 | "github.com/UnderTreeTech/waterdrop/pkg/trace" 32 | "github.com/opentracing/opentracing-go/log" 33 | 34 | "google.golang.org/grpc" 35 | ) 36 | 37 | // TraceForUnaryServer trace unary server side details 38 | func TraceForUnaryServer() grpc.UnaryServerInterceptor { 39 | return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) { 40 | opt := trace.FromIncomingContext(ctx) 41 | span, ctx := trace.StartSpanFromContext(ctx, info.FullMethod, opt) 42 | ext.Component.Set(span, "grpc") 43 | ext.SpanKind.Set(span, ext.SpanKindRPCServerEnum) 44 | if peer, ok := peer.FromContext(ctx); ok { 45 | ext.PeerAddress.Set(span, peer.Addr.String()) 46 | } 47 | defer span.Finish() 48 | 49 | resp, err = handler(ctx, req) 50 | if err != nil { 51 | estatus := status.ExtractStatus(err) 52 | ext.Error.Set(span, true) 53 | span.LogFields(log.String("event", "error"), log.Int("code", estatus.Code()), log.String("message", estatus.Message())) 54 | } 55 | 56 | return 57 | } 58 | } 59 | 60 | // TraceForUnaryClient trace unary client side details 61 | func TraceForUnaryClient() grpc.UnaryClientInterceptor { 62 | return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) { 63 | md, ok := metadata.FromOutgoingContext(ctx) 64 | if !ok { 65 | md = metadata.New(nil) 66 | } 67 | span, ctx := trace.StartSpanFromContext(ctx, method) 68 | ext.Component.Set(span, "grpc") 69 | ext.SpanKind.Set(span, ext.SpanKindRPCClientEnum) 70 | defer span.Finish() 71 | 72 | ctx = trace.MetadataInjector(ctx, md) 73 | err = invoker(ctx, method, req, reply, cc, opts...) 74 | if err != nil { 75 | estatus := status.ExtractStatus(err) 76 | ext.Error.Set(span, true) 77 | span.LogFields(log.String("event", "error"), log.Int("code", estatus.Code()), log.String("message", estatus.Message())) 78 | } 79 | return 80 | } 81 | } 82 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/trace_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "testing" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/trace" 27 | 28 | "github.com/opentracing/opentracing-go" 29 | jconfig "github.com/uber/jaeger-client-go/config" 30 | 31 | "github.com/UnderTreeTech/waterdrop/pkg/log" 32 | 33 | "github.com/stretchr/testify/assert" 34 | 35 | "google.golang.org/grpc" 36 | ) 37 | 38 | func newJaegerClient() (opentracing.Tracer, func()) { 39 | var configuration = jconfig.Configuration{ 40 | ServiceName: "trace", 41 | } 42 | 43 | tracer, closer, err := configuration.NewTracer() 44 | if err != nil { 45 | panic(fmt.Sprintf("new jaeger trace fail, err msg %s", err.Error())) 46 | } 47 | 48 | return tracer, func() { closer.Close() } 49 | } 50 | 51 | func TestTraceForUnaryServer(t *testing.T) { 52 | interceptor := TraceForUnaryServer() 53 | handler := func(ctx context.Context, req interface{}) (resp interface{}, err error) { 54 | log.Info(ctx, "test server trace") 55 | return 56 | } 57 | info := &grpc.UnaryServerInfo{ 58 | FullMethod: "/grpc.testing.TestService/UnaryCall", 59 | } 60 | 61 | t.Run("mock trace", func(t *testing.T) { 62 | resp, err := interceptor(context.Background(), nil, info, handler) 63 | assert.Nil(t, resp) 64 | assert.Nil(t, err) 65 | }) 66 | 67 | tracer, close := newJaegerClient() 68 | trace.SetGlobalTracer(tracer) 69 | defer close() 70 | t.Run("jaeger trace", func(t *testing.T) { 71 | resp, err := interceptor(context.Background(), nil, info, handler) 72 | assert.Nil(t, resp) 73 | assert.Nil(t, err) 74 | }) 75 | } 76 | 77 | func TestTraceForUnaryClient(t *testing.T) { 78 | interceptor := TraceForUnaryClient() 79 | invoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) (err error) { 80 | log.Info(ctx, "test client trace") 81 | return 82 | } 83 | 84 | t.Run("mock trace", func(t *testing.T) { 85 | err := interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 86 | assert.Nil(t, err) 87 | }) 88 | 89 | tracer, close := newJaegerClient() 90 | trace.SetGlobalTracer(tracer) 91 | defer close() 92 | t.Run("jaeger trace", func(t *testing.T) { 93 | err := interceptor(context.Background(), "/grpc.testing.TestService/UnaryCall", nil, nil, nil, invoker) 94 | assert.Nil(t, err) 95 | }) 96 | } 97 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/validator.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | 24 | "github.com/go-playground/validator/v10" 25 | 26 | "google.golang.org/grpc" 27 | ) 28 | 29 | var v = validator.New() 30 | 31 | // ValidateForUnaryServer validate input request params 32 | func ValidateForUnaryServer() grpc.UnaryServerInterceptor { 33 | return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) { 34 | if err = v.Struct(req); err != nil { 35 | return 36 | } 37 | return handler(ctx, req) 38 | } 39 | } 40 | 41 | // GetValidator returns the underlying validator engine which powers the 42 | // StructValidator implementation. 43 | func GetValidator() *validator.Validate { 44 | return v 45 | } 46 | -------------------------------------------------------------------------------- /pkg/server/rpc/interceptors/validator_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package interceptors 20 | 21 | import ( 22 | "context" 23 | "testing" 24 | 25 | "github.com/go-playground/validator/v10" 26 | 27 | "github.com/stretchr/testify/assert" 28 | "google.golang.org/grpc" 29 | ) 30 | 31 | type Mock struct { 32 | Name string `validate:"required,min=6,max=10"` 33 | Email string `validate:"required,email"` 34 | Age int32 `validate:"required,gte=1,lte=60"` 35 | } 36 | 37 | func TestValidateForUnaryServer(t *testing.T) { 38 | interceptor := ValidateForUnaryServer() 39 | handler := func(ctx context.Context, req interface{}) (resp interface{}, err error) { 40 | return nil, nil 41 | } 42 | info := &grpc.UnaryServerInfo{ 43 | FullMethod: "/grpc.testing.TestService/UnaryCall", 44 | } 45 | 46 | successMock := &Mock{ 47 | Name: "waterdrop", 48 | Email: "example@example.com", 49 | Age: 30, 50 | } 51 | t.Run("validate success", func(t *testing.T) { 52 | resp, err := interceptor(context.Background(), successMock, info, handler) 53 | assert.Nil(t, resp) 54 | assert.Nil(t, err) 55 | }) 56 | 57 | failMock := &Mock{ 58 | Name: "hello", 59 | Email: "waterdrop", 60 | Age: 100, 61 | } 62 | t.Run("validate fail", func(t *testing.T) { 63 | resp, err := interceptor(context.Background(), failMock, info, handler) 64 | assert.Nil(t, resp) 65 | assert.NotEqual(t, nil, err) 66 | }) 67 | } 68 | 69 | func TestGetValidator(t *testing.T) { 70 | v := GetValidator() 71 | assert.IsType(t, &validator.Validate{}, v) 72 | } 73 | -------------------------------------------------------------------------------- /pkg/server/rpc/metadata/metadata.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metadata 20 | 21 | import "math" 22 | 23 | const ( 24 | // MaxInterceptors max interceptors for rpc client/server 25 | // The idea borrowed from gin 26 | MaxInterceptors = math.MaxInt8 / 2 27 | ) 28 | -------------------------------------------------------------------------------- /pkg/server/rpc/server/server_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package server 20 | 21 | import ( 22 | "context" 23 | "os" 24 | "testing" 25 | "time" 26 | 27 | "github.com/UnderTreeTech/waterdrop/pkg/server/rpc/config" 28 | 29 | "github.com/stretchr/testify/assert" 30 | "google.golang.org/grpc/test/grpc_testing" 31 | 32 | "github.com/UnderTreeTech/waterdrop/pkg/log" 33 | ) 34 | 35 | var srv = New(config.DefaultServerConfig()) 36 | 37 | func TestMain(m *testing.M) { 38 | defer log.New(nil).Sync() 39 | 40 | code := m.Run() 41 | os.Exit(code) 42 | } 43 | 44 | func TestStart(t *testing.T) { 45 | grpc_testing.RegisterTestServiceServer(srv.server, &grpc_testing.UnimplementedTestServiceServer{}) 46 | net := srv.Start() 47 | assert.Equal(t, "[::]:20812", net.String()) 48 | assert.Equal(t, "tcp", net.Network()) 49 | } 50 | 51 | func TestStop(t *testing.T) { 52 | go func() { 53 | time.Sleep(100 * time.Millisecond) 54 | err := srv.Stop(context.Background()) 55 | assert.Equal(t, err, nil) 56 | }() 57 | 58 | time.Sleep(200 * time.Millisecond) 59 | } 60 | -------------------------------------------------------------------------------- /pkg/stats/metric/counter.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metric 20 | 21 | import "github.com/prometheus/client_golang/prometheus" 22 | 23 | // CounterVecOpts 24 | type CounterVecOpts struct { 25 | Namespace string 26 | Subsystem string 27 | Name string 28 | Help string 29 | Labels []string 30 | } 31 | 32 | type counterVec struct { 33 | *prometheus.CounterVec 34 | } 35 | 36 | // NewCounterVec returns a CounterVecOpts instance 37 | func NewCounterVec(opt *CounterVecOpts) *counterVec { 38 | vector := prometheus.NewCounterVec( 39 | prometheus.CounterOpts{ 40 | Namespace: opt.Namespace, 41 | Subsystem: opt.Subsystem, 42 | Name: opt.Name, 43 | Help: opt.Help, 44 | }, opt.Labels) 45 | prometheus.MustRegister(vector) 46 | 47 | return &counterVec{CounterVec: vector} 48 | } 49 | 50 | // Inc increments the counter by 1 51 | func (c *counterVec) Inc(labels ...string) { 52 | c.WithLabelValues(labels...).Inc() 53 | } 54 | 55 | // Add adds the given value to the counter. It panics if the value is < 0 56 | func (c *counterVec) Add(v float64, labels ...string) { 57 | c.WithLabelValues(labels...).Add(v) 58 | } 59 | -------------------------------------------------------------------------------- /pkg/stats/metric/gauge.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metric 20 | 21 | import "github.com/prometheus/client_golang/prometheus" 22 | 23 | // GaugeVecOpts 24 | type GaugeVecOpts struct { 25 | Namespace string 26 | Subsystem string 27 | Name string 28 | Help string 29 | Labels []string 30 | } 31 | 32 | type gaugeVec struct { 33 | *prometheus.GaugeVec 34 | } 35 | 36 | // NewGaugeVec returns a GaugeVecOpts instance 37 | func NewGaugeVec(opt *GaugeVecOpts) *gaugeVec { 38 | vector := prometheus.NewGaugeVec( 39 | prometheus.GaugeOpts{ 40 | Namespace: opt.Namespace, 41 | Subsystem: opt.Subsystem, 42 | Name: opt.Name, 43 | Help: opt.Help, 44 | }, opt.Labels) 45 | prometheus.MustRegister(vector) 46 | 47 | return &gaugeVec{GaugeVec: vector} 48 | } 49 | 50 | // Inc increments the Gauge by 1 51 | func (g *gaugeVec) Inc(labels ...string) { 52 | g.WithLabelValues(labels...).Inc() 53 | } 54 | 55 | // Dec decrements the Gauge by 1 56 | func (g *gaugeVec) Dec(labels ...string) { 57 | g.WithLabelValues(labels...).Dec() 58 | } 59 | 60 | // Add adds the given value to the Gauge 61 | func (g *gaugeVec) Add(v float64, labels ...string) { 62 | g.WithLabelValues(labels...).Add(v) 63 | } 64 | 65 | // Sub subtracts the given value from the Gauge 66 | func (g *gaugeVec) Sub(v float64, labels ...string) { 67 | g.WithLabelValues(labels...).Sub(v) 68 | } 69 | -------------------------------------------------------------------------------- /pkg/stats/metric/histogram.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metric 20 | 21 | import "github.com/prometheus/client_golang/prometheus" 22 | 23 | // HistogramVecOpts 24 | type HistogramVecOpts struct { 25 | Namespace string 26 | Subsystem string 27 | Name string 28 | Help string 29 | Labels []string 30 | Buckets []float64 31 | } 32 | 33 | type histogramVec struct { 34 | *prometheus.HistogramVec 35 | } 36 | 37 | // NewHistogramVec returns a HistogramVecOpts instance 38 | func NewHistogramVec(opt *HistogramVecOpts) *histogramVec { 39 | vector := prometheus.NewHistogramVec( 40 | prometheus.HistogramOpts{ 41 | Namespace: opt.Namespace, 42 | Subsystem: opt.Subsystem, 43 | Name: opt.Name, 44 | Help: opt.Help, 45 | Buckets: opt.Buckets, 46 | }, opt.Labels) 47 | prometheus.MustRegister(vector) 48 | 49 | return &histogramVec{HistogramVec: vector} 50 | } 51 | 52 | // Observe add observations to Histogram. 53 | func (h *histogramVec) Observe(v float64, labels ...string) { 54 | h.WithLabelValues(labels...).Observe(v) 55 | } 56 | -------------------------------------------------------------------------------- /pkg/stats/metric/metric.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package metric 20 | 21 | import ( 22 | "net/http" 23 | 24 | "github.com/gin-gonic/gin" 25 | "github.com/prometheus/client_golang/prometheus/promhttp" 26 | ) 27 | 28 | // RegisterMetric register metric handler 29 | func RegisterMetric(engine *gin.Engine) { 30 | metric := engine.Group("/metrics") 31 | { 32 | metric.GET("/", metricHandler(promhttp.Handler().ServeHTTP)) 33 | } 34 | } 35 | 36 | // metricHandler metric handler 37 | func metricHandler(handler http.HandlerFunc) gin.HandlerFunc { 38 | return func(ctx *gin.Context) { 39 | handler.ServeHTTP(ctx.Writer, ctx.Request) 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /pkg/stats/profile/profile.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package profile 20 | 21 | import ( 22 | "net/http" 23 | "net/http/pprof" 24 | 25 | "github.com/gin-gonic/gin" 26 | ) 27 | 28 | // RegisterProfile register profile handler 29 | func RegisterProfile(engine *gin.Engine) { 30 | perf := engine.Group("/debug/profile") 31 | { 32 | perf.GET("/", profileHandler(pprof.Index)) 33 | perf.GET("/cmdline", profileHandler(pprof.Cmdline)) 34 | perf.GET("/profile", profileHandler(pprof.Profile)) 35 | perf.GET("/symbol", profileHandler(pprof.Symbol)) 36 | perf.GET("/trace", profileHandler(pprof.Trace)) 37 | perf.GET("/allocs", profileHandler(pprof.Handler("allocs").ServeHTTP)) 38 | perf.GET("/block", profileHandler(pprof.Handler("block").ServeHTTP)) 39 | perf.GET("/goroutine", profileHandler(pprof.Handler("goroutine").ServeHTTP)) 40 | perf.GET("/heap", profileHandler(pprof.Handler("heap").ServeHTTP)) 41 | perf.GET("/mutex", profileHandler(pprof.Handler("mutex").ServeHTTP)) 42 | perf.GET("/threadcreate", profileHandler(pprof.Handler("threadcreate").ServeHTTP)) 43 | } 44 | } 45 | 46 | // profileHandler performance handler 47 | func profileHandler(handler http.HandlerFunc) gin.HandlerFunc { 48 | return func(ctx *gin.Context) { 49 | handler.ServeHTTP(ctx.Writer, ctx.Request) 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /pkg/stats/stats.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package stats 20 | 21 | import ( 22 | "fmt" 23 | "net" 24 | 25 | "github.com/UnderTreeTech/waterdrop/pkg/registry" 26 | "github.com/UnderTreeTech/waterdrop/pkg/utils/xnet" 27 | 28 | "github.com/gin-gonic/gin" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/stats/metric" 31 | "github.com/UnderTreeTech/waterdrop/pkg/stats/profile" 32 | ) 33 | 34 | // StartStats start stats server 35 | func StartStats(serviceName string) (si *registry.ServiceInfo, err error) { 36 | gin.SetMode("release") 37 | engine := gin.Default() 38 | 39 | profile.RegisterProfile(engine) 40 | metric.RegisterMetric(engine) 41 | listener, err := net.Listen("tcp", ":0") 42 | if err != nil { 43 | return nil, err 44 | } 45 | 46 | go func() { 47 | if err = engine.RunListener(listener); err != nil { 48 | return 49 | } 50 | }() 51 | 52 | _, port, _ := net.SplitHostPort(listener.Addr().String()) 53 | si = ®istry.ServiceInfo{ 54 | Name: serviceName + "/stats", 55 | Scheme: "http", 56 | Addr: fmt.Sprintf("%s://%s:%s", "http", xnet.InternalIP(), port), 57 | Version: "1.0.0", 58 | } 59 | return 60 | } 61 | -------------------------------------------------------------------------------- /pkg/status/README.md: -------------------------------------------------------------------------------- 1 | Status is used to unified errors between services, no matter is http/rpc. 2 | 3 | The code is 60% copied from the official gRPC status package. -------------------------------------------------------------------------------- /pkg/store/minio/minio_test.go: -------------------------------------------------------------------------------- 1 | package minio 2 | 3 | import ( 4 | "context" 5 | "os" 6 | "testing" 7 | 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestGetFileUrl(t *testing.T) { 12 | cfg := &Config{ 13 | InternalEndpoint: "localhost:9000", 14 | ExternalEndpoint: "localhost:9000", 15 | Region: "", 16 | AccessKey: "XMFMOKB2FJWA0I9JIR62", 17 | SecretKey: "aMRKmxoRcb+Ezr5CmOmFAqFwYWPrEFA7UdtWWWOl", 18 | } 19 | clnt, err := New(cfg) 20 | assert.Nil(t, err) 21 | ctx := context.Background() 22 | 23 | exsit, err := clnt.ExistBucket(ctx, "test") 24 | assert.Nil(t, err) 25 | if !exsit { 26 | err = clnt.CreateBucket(ctx, "test") 27 | assert.Nil(t, err) 28 | } 29 | 30 | _, err = clnt.FPutObject(ctx, "test", "minio.go", "./minio.go") 31 | assert.Nil(t, err) 32 | 33 | object, err := os.Open("./minio_test.go") 34 | assert.Nil(t, err) 35 | defer object.Close() 36 | 37 | stat, err := object.Stat() 38 | assert.Nil(t, err) 39 | _, err = clnt.PutObject(ctx, "test", "minio_test.go", object, stat.Size(), ".jpeg") 40 | assert.Nil(t, err) 41 | path, err := clnt.GetFileUrl(ctx, "test", "minio_test.go") 42 | assert.Nil(t, err) 43 | assert.Contains(t, path, "minio_test.go") 44 | } 45 | -------------------------------------------------------------------------------- /pkg/trace/metadata.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package trace 20 | 21 | import "strings" 22 | 23 | type CarrierMD struct { 24 | md map[string][]string 25 | } 26 | 27 | // Set a key:value pair to the carrier. Multiple calls to Set() for the 28 | // same key leads to undefined behavior. 29 | func (cm CarrierMD) Set(key, val string) { 30 | key = strings.ToLower(key) 31 | cm.md[key] = append(cm.md[key], val) 32 | } 33 | 34 | // ForeachKey returns TextMap contents via repeated calls to the `handler` 35 | // function. If any call to `handler` returns a non-nil error, ForeachKey 36 | // terminates and returns that error. 37 | func (cm CarrierMD) ForeachKey(handler func(key, val string) error) error { 38 | for k, vals := range cm.md { 39 | for _, v := range vals { 40 | if err := handler(k, v); err != nil { 41 | return err 42 | } 43 | } 44 | } 45 | return nil 46 | } 47 | -------------------------------------------------------------------------------- /pkg/utils/xbuffer/README.md: -------------------------------------------------------------------------------- 1 | ## xbuffer 2 | 3 | 1. BufferPool 4 | Based on sync.Pool, sync.Pool automatically GC's objects. 5 | 6 | 2. SizedBufferPool 7 | Based channel, creates a new BufferPool bounded to the given size. objects in this pool are 8 | alive along with the application life cycle. -------------------------------------------------------------------------------- /pkg/utils/xbuffer/bufferpool.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xbuffer 20 | 21 | import ( 22 | "bytes" 23 | "sync" 24 | ) 25 | 26 | // BufferPool buffer pool definition 27 | type BufferPool struct { 28 | //alloc sets the initial capacity of new buffers to minimize calls to make() 29 | alloc int 30 | pool *sync.Pool 31 | } 32 | 33 | // NewBufferPool creates a new BufferPool bounded to the given buffer size. 34 | func NewBufferPool(alloc int) *BufferPool { 35 | return &BufferPool{ 36 | alloc: alloc, 37 | pool: &sync.Pool{ 38 | New: func() interface{} { 39 | return bytes.NewBuffer(make([]byte, 0, alloc)) 40 | }, 41 | }, 42 | } 43 | } 44 | 45 | // Get gets a Buffer from the BufferPool 46 | func (bp *BufferPool) Get() *bytes.Buffer { 47 | buffer := bp.pool.Get().(*bytes.Buffer) 48 | return buffer 49 | } 50 | 51 | // Put returns the given Buffer to the BufferPool 52 | func (bp *BufferPool) Put(buffer *bytes.Buffer) { 53 | if buffer.Cap() <= bp.alloc { 54 | buffer.Reset() 55 | bp.pool.Put(buffer) 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /pkg/utils/xbuffer/bufferpool_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xbuffer 20 | 21 | import ( 22 | "bytes" 23 | "testing" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | func TestBufferPool(t *testing.T) { 29 | alloc := 1024 30 | pool := NewBufferPool(alloc) 31 | pool.Put(bytes.NewBuffer(make([]byte, 0, 2*alloc))) 32 | assert.True(t, pool.Get().Cap() <= alloc) 33 | } 34 | -------------------------------------------------------------------------------- /pkg/utils/xbuffer/sizedbufferpool.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xbuffer 20 | 21 | import "bytes" 22 | 23 | // SizedBufferPool sized buffer pool definition 24 | type SizedBufferPool struct { 25 | pool chan *bytes.Buffer 26 | size int 27 | alloc int 28 | } 29 | 30 | // SizedBufferPool creates a new BufferPool bounded to the given size. 31 | // size defines the number of buffers to be retained in the pool and alloc sets 32 | // the initial capacity of new buffers to minimize calls to make(). 33 | // 34 | // The value of alloc should seek to provide a buffer that is representative of 35 | // most data written to the the buffer (i.e. 95th percentile) without being 36 | // overly large (which will increase static memory consumption). 37 | func NewSizedBufferPool(size int, alloc int) (bp *SizedBufferPool) { 38 | return &SizedBufferPool{ 39 | size: size, 40 | pool: make(chan *bytes.Buffer, size), 41 | alloc: alloc, 42 | } 43 | } 44 | 45 | // Get gets a Buffer from the SizedBufferPool, or creates a new one if none are 46 | // available in the pool. Buffers have a pre-allocated capacity. 47 | func (bp *SizedBufferPool) Get() (b *bytes.Buffer) { 48 | select { 49 | case b = <-bp.pool: 50 | // reuse existing buffer 51 | default: 52 | // create new buffer 53 | b = bytes.NewBuffer(make([]byte, 0, bp.alloc)) 54 | } 55 | return 56 | } 57 | 58 | // Put returns the given Buffer to the SizedBufferPool. 59 | func (bp *SizedBufferPool) Put(b *bytes.Buffer) { 60 | if b.Cap() > bp.alloc { 61 | if len(bp.pool) < bp.size { 62 | b = bytes.NewBuffer(make([]byte, 0, bp.alloc)) 63 | } else { 64 | return 65 | } 66 | } else { 67 | b.Reset() 68 | } 69 | 70 | select { 71 | case bp.pool <- b: 72 | default: // Discard the buffer if the pool is full. 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /pkg/utils/xbuffer/sizedbufferpool_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xbuffer 20 | 21 | import ( 22 | "bytes" 23 | "testing" 24 | ) 25 | 26 | // TestSizedBufferPool checks that over-sized buffers are released and that new 27 | // buffers are created in their place. 28 | func TestSizedBufferPool(t *testing.T) { 29 | size := 4 30 | capacity := 1024 31 | 32 | bufPool := NewSizedBufferPool(size, capacity) 33 | 34 | b := bufPool.Get() 35 | 36 | // Check the cap before we use the buffer. 37 | if cap(b.Bytes()) != capacity { 38 | t.Fatalf("buffer capacity incorrect: got %v want %v", cap(b.Bytes()), 39 | capacity) 40 | } 41 | 42 | // Grow the buffer beyond our capacity and return it to the pool 43 | b.Grow(capacity * 3) 44 | bufPool.Put(b) 45 | 46 | // Add some additional buffers to fill up the pool. 47 | for i := 0; i < size; i++ { 48 | bufPool.Put(bytes.NewBuffer(make([]byte, 0, bufPool.alloc*2))) 49 | } 50 | 51 | // Check that oversized buffers are being replaced. 52 | if len(bufPool.pool) < size { 53 | t.Fatalf("buffer pool too small: got %v want %v", len(bufPool.pool), size) 54 | } 55 | 56 | // Close the channel so we can iterate over it. 57 | close(bufPool.pool) 58 | 59 | // Check that there are buffers of the correct capacity in the pool. 60 | for buffer := range bufPool.pool { 61 | if cap(buffer.Bytes()) != bufPool.alloc { 62 | t.Fatalf("returned buffers wrong capacity: got %v want %v", 63 | cap(buffer.Bytes()), capacity) 64 | } 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /pkg/utils/xcollection/lru_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcollection 20 | 21 | import ( 22 | "fmt" 23 | "testing" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | type simpleStruct struct { 29 | int 30 | string 31 | } 32 | 33 | type complexStruct struct { 34 | int 35 | simpleStruct 36 | } 37 | 38 | var getTests = []struct { 39 | name string 40 | keyToAdd interface{} 41 | keyToGet interface{} 42 | expectedOk bool 43 | }{ 44 | {"string_hit", "myKey", "myKey", true}, 45 | {"string_miss", "myKey", "nonsense", false}, 46 | {"simple_struct_hit", simpleStruct{1, "two"}, simpleStruct{1, "two"}, true}, 47 | {"simple_struct_miss", simpleStruct{1, "two"}, simpleStruct{0, "noway"}, false}, 48 | {"complex_struct_hit", complexStruct{1, simpleStruct{2, "three"}}, 49 | complexStruct{1, simpleStruct{2, "three"}}, true}, 50 | } 51 | 52 | func TestGet(t *testing.T) { 53 | lru := NewLRU(0) 54 | for _, tt := range getTests { 55 | lru.Add(tt.keyToAdd, 1234) 56 | _, ok := lru.Get(tt.keyToGet) 57 | assert.Equal(t, ok, tt.expectedOk) 58 | } 59 | } 60 | 61 | func TestRemove(t *testing.T) { 62 | lru := NewLRU(0) 63 | lru.Add("myKey", 1234) 64 | val, ok := lru.Get("myKey") 65 | assert.Equal(t, true, ok) 66 | assert.Equal(t, 1234, val.(int)) 67 | 68 | lru.Remove("myKey") 69 | _, ok = lru.Get("myKey") 70 | assert.Equal(t, false, ok) 71 | } 72 | 73 | func TestEvict(t *testing.T) { 74 | evictedKeys := make([]Key, 0) 75 | onEvictedFun := func(key Key, value interface{}) { 76 | evictedKeys = append(evictedKeys, key) 77 | } 78 | 79 | lru := NewLRU(20) 80 | lru.OnEvicted = onEvictedFun 81 | for i := 0; i < 22; i++ { 82 | lru.Add(fmt.Sprintf("myKey%d", i), 1234) 83 | } 84 | 85 | assert.Equal(t, 2, len(evictedKeys)) 86 | assert.Equal(t, evictedKeys[0], Key("myKey0")) 87 | assert.Equal(t, evictedKeys[1], Key("myKey1")) 88 | } 89 | 90 | func TestClear(t *testing.T) { 91 | lru := NewLRU(20) 92 | for i := 0; i < 22; i++ { 93 | lru.Add(fmt.Sprintf("myKey%d", i), 1234) 94 | } 95 | lru.Clear() 96 | assert.Nil(t, lru.ll) 97 | assert.Nil(t, lru.cache) 98 | } 99 | 100 | func TestLRULen(t *testing.T) { 101 | lru := NewLRU(20) 102 | for i := 0; i < 22; i++ { 103 | lru.Add(fmt.Sprintf("myKey%d", i), 1234) 104 | } 105 | assert.Equal(t, lru.Len(), 20) 106 | } 107 | -------------------------------------------------------------------------------- /pkg/utils/xcollection/priorityqueue.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcollection 20 | 21 | import "container/heap" 22 | 23 | type Item struct { 24 | Value interface{} 25 | Priority int64 26 | Index int 27 | } 28 | 29 | // this is a priority queue as implemented by a min heap 30 | // ie. the 0th element is the *lowest* value 31 | type PriorityQueue []*Item 32 | 33 | func NewPriorityQueue(capacity int) PriorityQueue { 34 | return make(PriorityQueue, 0, capacity) 35 | } 36 | 37 | // Len returns the number of items in the queue 38 | func (pq PriorityQueue) Len() int { 39 | return len(pq) 40 | } 41 | 42 | // Less check the priority of two element in the queue 43 | func (pq PriorityQueue) Less(i, j int) bool { 44 | return pq[i].Priority < pq[j].Priority 45 | } 46 | 47 | // Swap swap two element index in the queue 48 | func (pq PriorityQueue) Swap(i, j int) { 49 | pq[i], pq[j] = pq[j], pq[i] 50 | pq[i].Index = i 51 | pq[j].Index = j 52 | } 53 | 54 | // Push push an element into the queue 55 | func (pq *PriorityQueue) Push(x interface{}) { 56 | n := len(*pq) 57 | c := cap(*pq) 58 | if n+1 > c { 59 | npq := make(PriorityQueue, n, c*2) 60 | copy(npq, *pq) 61 | *pq = npq 62 | } 63 | *pq = (*pq)[0 : n+1] 64 | item := x.(*Item) 65 | item.Index = n 66 | (*pq)[n] = item 67 | } 68 | 69 | // Pop pop out an element from the queue 70 | func (pq *PriorityQueue) Pop() interface{} { 71 | n := len(*pq) 72 | c := cap(*pq) 73 | if n < (c/2) && c > 25 { 74 | npq := make(PriorityQueue, n, c/2) 75 | copy(npq, *pq) 76 | *pq = npq 77 | } 78 | item := (*pq)[n-1] 79 | item.Index = -1 80 | *pq = (*pq)[0 : n-1] 81 | return item 82 | } 83 | 84 | // PeekAndShift return the first element or nil, and update the item priority at the same time 85 | func (pq *PriorityQueue) PeekAndShift(max int64) (*Item, int64) { 86 | if pq.Len() == 0 { 87 | return nil, 0 88 | } 89 | 90 | item := (*pq)[0] 91 | if item.Priority > max { 92 | return nil, item.Priority - max 93 | } 94 | heap.Remove(pq, 0) 95 | 96 | return item, 0 97 | } 98 | -------------------------------------------------------------------------------- /pkg/utils/xcollection/priorityqueue_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcollection 20 | 21 | import ( 22 | "container/heap" 23 | "math/rand" 24 | "sort" 25 | "testing" 26 | 27 | "github.com/stretchr/testify/assert" 28 | ) 29 | 30 | func TestNewPriorityQueue(t *testing.T) { 31 | c := 100 32 | pq := NewPriorityQueue(c) 33 | 34 | for i := 0; i < c+1; i++ { 35 | heap.Push(&pq, &Item{Value: i, Priority: int64(i)}) 36 | } 37 | assert.Equal(t, pq.Len(), c+1) 38 | assert.Equal(t, cap(pq), c*2) 39 | 40 | for i := 0; i < c+1; i++ { 41 | item := heap.Pop(&pq) 42 | assert.Equal(t, item.(*Item).Value.(int), i) 43 | } 44 | assert.Equal(t, cap(pq), c/4) 45 | } 46 | 47 | func TestPush(t *testing.T) { 48 | c := 100 49 | pq := NewPriorityQueue(c) 50 | ints := make([]int, 0, c) 51 | 52 | for i := 0; i < c; i++ { 53 | v := rand.Int() 54 | ints = append(ints, v) 55 | heap.Push(&pq, &Item{Value: i, Priority: int64(v)}) 56 | } 57 | assert.Equal(t, pq.Len(), c) 58 | assert.Equal(t, cap(pq), c) 59 | 60 | sort.Ints(ints) 61 | 62 | for i := 0; i < c; i++ { 63 | item, _ := pq.PeekAndShift(int64(ints[len(ints)-1])) 64 | assert.Equal(t, item.Priority, int64(ints[i])) 65 | } 66 | } 67 | 68 | func TestPop(t *testing.T) { 69 | c := 100 70 | pq := NewPriorityQueue(c) 71 | 72 | for i := 0; i < c; i++ { 73 | v := rand.Int() 74 | heap.Push(&pq, &Item{Value: "test", Priority: int64(v)}) 75 | } 76 | 77 | for i := 0; i < 10; i++ { 78 | heap.Remove(&pq, rand.Intn((c-1)-i)) 79 | } 80 | 81 | lastPriority := heap.Pop(&pq).(*Item).Priority 82 | for i := 0; i < (c - 10 - 1); i++ { 83 | item := heap.Pop(&pq) 84 | assert.Equal(t, lastPriority < item.(*Item).Priority, true) 85 | lastPriority = item.(*Item).Priority 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /pkg/utils/xcollection/rollingwindow_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcollection 20 | 21 | import ( 22 | "testing" 23 | "time" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | func TestRollingWindow(t *testing.T) { 29 | bucketSize := 10 30 | bucketDuration := time.Millisecond * 100 31 | window := NewRollingWindow(bucketSize, bucketDuration) 32 | 33 | for i := 0; i < bucketSize; i++ { 34 | window.Add(1) 35 | time.Sleep(bucketDuration) 36 | } 37 | 38 | var total int64 39 | var success float64 40 | window.Reduce(func(bucket *Bucket) { 41 | success += bucket.Sum 42 | total += bucket.Count 43 | }) 44 | 45 | assert.Equal(t, int64(success), total) 46 | } 47 | -------------------------------------------------------------------------------- /pkg/utils/xcollection/safemap_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcollection 20 | 21 | import ( 22 | "math/rand" 23 | "runtime" 24 | "sync" 25 | "testing" 26 | ) 27 | 28 | func TestConcurrentRange(t *testing.T) { 29 | const mapSize = 1 << 10 30 | 31 | m := NewSafeMap() 32 | for n := int64(1); n <= mapSize; n++ { 33 | m.Store(n, n) 34 | } 35 | 36 | done := make(chan struct{}) 37 | var wg sync.WaitGroup 38 | defer func() { 39 | close(done) 40 | wg.Wait() 41 | }() 42 | for g := int64(runtime.GOMAXPROCS(0)); g > 0; g-- { 43 | r := rand.New(rand.NewSource(g)) 44 | wg.Add(1) 45 | go func(g int64) { 46 | defer wg.Done() 47 | for i := int64(0); ; i++ { 48 | select { 49 | case <-done: 50 | return 51 | default: 52 | } 53 | for n := int64(1); n < mapSize; n++ { 54 | if r.Int63n(mapSize) == 0 { 55 | m.Store(n, n*i*g) 56 | } else { 57 | m.Load(n) 58 | } 59 | } 60 | } 61 | }(g) 62 | } 63 | 64 | iters := 1 << 10 65 | if testing.Short() { 66 | iters = 16 67 | } 68 | for n := iters; n > 0; n-- { 69 | seen := make(map[int64]bool, mapSize) 70 | 71 | m.Range(func(ki, vi interface{}) bool { 72 | k, v := ki.(int64), vi.(int64) 73 | if v%k != 0 { 74 | t.Fatalf("while Storing multiples of %v, Range saw value %v", k, v) 75 | } 76 | if seen[k] { 77 | t.Fatalf("Range visited key %v twice", k) 78 | } 79 | seen[k] = true 80 | return true 81 | }) 82 | 83 | if len(seen) != mapSize { 84 | t.Fatalf("Range visited %v elements of %v-element Map", len(seen), mapSize) 85 | } 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /pkg/utils/xcrypto/aes_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcrypto 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestAesCbc(t *testing.T) { 28 | content := []byte("Hello Waterdrop!") 29 | cbc, err := NewAesCbcCrypto("abcdefghijklmnop") 30 | assert.Equal(t, nil, err) 31 | encypt, err := cbc.Encrypt(content) 32 | assert.Nil(t, err) 33 | decrypt, err := cbc.Decrypt(encypt) 34 | assert.Nil(t, err) 35 | assert.Equal(t, content, decrypt) 36 | baseEncrypt, err := cbc.EncryptToString(content, Base64) 37 | assert.Nil(t, err) 38 | dst, err := cbc.DecryptFromString(baseEncrypt, Base64) 39 | assert.Equal(t, nil, err) 40 | assert.Equal(t, content, dst) 41 | } 42 | -------------------------------------------------------------------------------- /pkg/utils/xcrypto/hash.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcrypto 20 | 21 | import ( 22 | "crypto/hmac" 23 | "crypto/md5" 24 | "crypto/sha1" 25 | "crypto/sha256" 26 | "crypto/sha512" 27 | "hash" 28 | "io" 29 | 30 | "github.com/UnderTreeTech/waterdrop/pkg/utils/xstring" 31 | ) 32 | 33 | // Hash hash input data using assigned hash mode 34 | func Hash(src string, mode HashMode) (dst []byte) { 35 | dst = hashEncrypt(xstring.StringToBytes(src), mode) 36 | return 37 | } 38 | 39 | // HashToString hash input to assigned encode mode. 40 | func HashToString(src string, mode HashMode, encode Encode) (dst string, err error) { 41 | encrypt := Hash(src, mode) 42 | dst, err = EncodeToString(encrypt, encode) 43 | return 44 | } 45 | 46 | // hashEncrypt encrypt input using assigned hash mode 47 | // hash mode: md5, sha1, sha256, sha512, default md5 48 | func hashEncrypt(data []byte, mode HashMode) (hashed []byte) { 49 | f := getHashFunc(mode) 50 | hh := f() 51 | hh.Write(data) 52 | hashed = hh.Sum(nil) 53 | return 54 | } 55 | 56 | // getHashFunc gets the crypto hash func 57 | func getHashFunc(mode HashMode) (f func() hash.Hash) { 58 | switch mode { 59 | case SHA1: 60 | return sha1.New 61 | case SHA256: 62 | return sha256.New 63 | case SHA512: 64 | return sha512.New 65 | case MD5: 66 | return md5.New 67 | default: 68 | return md5.New 69 | } 70 | } 71 | 72 | // HmacSHA256 returns HMAC bytes for src with the given key 73 | func HmacSHA256(key []byte, src string) (dst []byte) { 74 | h := hmac.New(getHashFunc(SHA256), key) 75 | io.WriteString(h, src) 76 | dst = h.Sum(nil) 77 | return 78 | } 79 | 80 | // HmacSHA256ToString returns hmac hash string with the given key and assigned encode mode 81 | // hash mode: md5, sha1, sha256, sha512, default md5 82 | func HmacSHA256ToString(key []byte, src string, encode Encode) (dst string, err error) { 83 | encrypt := HmacSHA256(key, src) 84 | dst, err = EncodeToString(encrypt, encode) 85 | return 86 | } 87 | -------------------------------------------------------------------------------- /pkg/utils/xcrypto/hash_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcrypto 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestHash(t *testing.T) { 28 | content := "123456" 29 | md5Str := "e10adc3949ba59abbe56e057f20f883e" 30 | sha1Str := "7c4a8d09ca3762af61e59520943dc26494f8941b" 31 | sha256Str := "8d969eef6ecad3c29a3a629280e686cf0c3f5d5a86aff3ca12020c923adc6c92" 32 | sha512Str := "ba3253876aed6bc22d4a6ff53d8406c6ad864195ed144ab5c87621b6c233b548baeae6956df346ec8c17f5ea10f35ee3cbc514797ed7ddd3145464e2a0bab413" 33 | 34 | encryptMd5Str, _ := HashToString(content, MD5, HEX) 35 | assert.Equal(t, md5Str, encryptMd5Str) 36 | encryptSha1Str, _ := HashToString(content, SHA1, HEX) 37 | assert.Equal(t, sha1Str, encryptSha1Str) 38 | encryptSha256Str, _ := HashToString(content, SHA256, HEX) 39 | assert.Equal(t, sha256Str, encryptSha256Str) 40 | encryptSha512Str, _ := HashToString(content, SHA512, HEX) 41 | assert.Equal(t, sha512Str, encryptSha512Str) 42 | } 43 | 44 | func TestHmacSHA256(t *testing.T) { 45 | content := "123456" 46 | key := "waterdrop" 47 | encrypt := "74a05d4cf0a06caafa4e3a8cbe3bb04a1476178aed8cb53dc109edd8ae3c3f30" 48 | 49 | dst, err := HmacSHA256ToString([]byte(key), content, HEX) 50 | assert.Nil(t, err) 51 | assert.Equal(t, encrypt, dst) 52 | } 53 | -------------------------------------------------------------------------------- /pkg/utils/xcrypto/pem/rsa_public_key.pem: -------------------------------------------------------------------------------- 1 | -----BEGIN RSA PUBLIC KEY----- 2 | MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA84ftEFkil1tLdA01+QTF 3 | 0AGDZ+k5vmT3J68XzH4Gysn/Aq3whqyfwF074dp7FXtpsVz0L6FrWzSq9jSjE8Nz 4 | ZXFBToTuhqXp04e2k9DtbJan1TZv9VVrVwxa6t8gvX9o90v2CJyxRMsuhQTzkhz4 5 | HjxJbnFV7SKLhIUAOyNg3hx0PduMHBjyjLpL90R4vIbvlIKVyzH3mJ0Vw4DRbcDN 6 | 7fXVaGVEYa+chHwwipcL2OZhrmjmIFVyrSkRMWw3DUdT2Or2l5lkD4dALz5JMuUP 7 | cpVWxouL1l1bWoOjBqTlK/LQhGI9i8v0227frjFXrClt/SZCj8xk73Rm/4BWIuPx 8 | y8ZsruSkxsUlEp04Z6dOFB25sF9532WtxEZeKWRdSL7oh8t+OZkeswpaYUg0UjV3 9 | /yLsP5a+sxjwakazH6irfQ0ak5xDemhsyNKR/o4b+97U8R9aa8V4AFDyHzViGHbg 10 | PfPEf/xyNDoMDfYLnXGauvzlqiCw4a8Z05K4ejtrbQQApsRWzN06/I2o2Jwq0K0q 11 | TkUCLnCF/PE7bhFooPCyqOQtIQ5t1C33or10k08gxJKmmRBYYrI7ztuxO8j5DdPb 12 | /JWWt49GPT9iPW6ChNrpiuimAP3LILOQhX+XQv8JN5g/G52AC0h8a7dDm5CWDdB9 13 | vsgqyu2A3IWy/QnceJxLVv8CAwEAAQ== 14 | -----END RSA PUBLIC KEY----- -------------------------------------------------------------------------------- /pkg/utils/xcrypto/rsa_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcrypto 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestRsa(t *testing.T) { 28 | content := "High performance go framework waterdrop" 29 | config := &RsaConfig{ 30 | PrivateKeyPath: "./pem/rsa_private_key.pem", 31 | PublicKeyPath: "./pem/rsa_public_key.pem", 32 | } 33 | rsa, err := NewRsaCrypt(config) 34 | assert.Equal(t, err, nil) 35 | 36 | encrypt, err := rsa.Encrypt([]byte(content)) 37 | assert.Equal(t, err, nil) 38 | decrypt, err := rsa.Decrypt(encrypt) 39 | assert.Equal(t, err, nil) 40 | assert.Equal(t, content, string(decrypt)) 41 | 42 | hexEncrypt, err := rsa.EncryptToString([]byte(content), HEX) 43 | assert.Equal(t, err, nil) 44 | hexDecrypt, err := rsa.DecryptFromString(hexEncrypt, HEX) 45 | assert.Equal(t, err, nil) 46 | assert.Equal(t, content, string(hexDecrypt)) 47 | 48 | base64Encrypt, err := rsa.EncryptToString([]byte(content), Base64) 49 | assert.Equal(t, err, nil) 50 | base64Decrypt, err := rsa.DecryptFromString(base64Encrypt, Base64) 51 | assert.Equal(t, err, nil) 52 | assert.Equal(t, content, string(base64Decrypt)) 53 | 54 | byteEncrypt, err := rsa.EncryptToString([]byte(content), PlainText) 55 | assert.Equal(t, err, nil) 56 | byteDecrypt, err := rsa.DecryptFromString(byteEncrypt, PlainText) 57 | assert.Equal(t, err, nil) 58 | assert.Equal(t, content, string(byteDecrypt)) 59 | } 60 | -------------------------------------------------------------------------------- /pkg/utils/xcrypto/util.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xcrypto 20 | 21 | import ( 22 | "encoding/base64" 23 | "encoding/hex" 24 | "errors" 25 | 26 | "github.com/UnderTreeTech/waterdrop/pkg/utils/xstring" 27 | ) 28 | 29 | //Hash for crypto Hash 30 | type HashMode uint 31 | 32 | const ( 33 | MD5 HashMode = iota 34 | SHA1 35 | SHA256 36 | SHA512 37 | ) 38 | 39 | //Encode defines the type of bytes encoded to string 40 | type Encode uint 41 | 42 | const ( 43 | PlainText Encode = iota // plain text, no encode just string 44 | HEX // hex encode 45 | Base64 // base64 encode 46 | ) 47 | 48 | // ErrInvalidEncodeMode unsupported encode mode 49 | var ErrInvalidEncodeMode = errors.New("unsupported encode mode") 50 | 51 | // DecodeString decodes string data to bytes in designed encoded type 52 | func DecodeString(data string, encodedType Encode) (dst []byte, err error) { 53 | switch encodedType { 54 | case PlainText: 55 | dst = xstring.StringToBytes(data) 56 | case HEX: 57 | dst, err = hex.DecodeString(data) 58 | case Base64: 59 | dst, err = base64.StdEncoding.DecodeString(data) 60 | default: 61 | return dst, ErrInvalidEncodeMode 62 | } 63 | return 64 | } 65 | 66 | //EncodeToString encodes data to string with encode type 67 | func EncodeToString(data []byte, encodeType Encode) (dst string, err error) { 68 | switch encodeType { 69 | case HEX: 70 | dst = hex.EncodeToString(data) 71 | case Base64: 72 | dst = base64.StdEncoding.EncodeToString(data) 73 | case PlainText: 74 | dst = xstring.BytesToString(data) 75 | default: 76 | return "", ErrInvalidEncodeMode 77 | } 78 | return 79 | } 80 | -------------------------------------------------------------------------------- /pkg/utils/xdefer/defers.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xdefer 20 | 21 | import "sync" 22 | 23 | // Defers wrap defer func 24 | type Defers struct { 25 | callbacks []func() error 26 | m sync.Mutex 27 | } 28 | 29 | // New return a Defers pointer 30 | func New() *Defers { 31 | return &Defers{ 32 | callbacks: make([]func() error, 0), 33 | } 34 | } 35 | 36 | // Add add a func to Defers 37 | func (d *Defers) Add(fns ...func() error) { 38 | d.m.Lock() 39 | defer d.m.Unlock() 40 | 41 | d.callbacks = append(d.callbacks, fns...) 42 | } 43 | 44 | // Close close defer funcs in Defers 45 | func (d *Defers) Close() { 46 | d.m.Lock() 47 | defer d.m.Unlock() 48 | 49 | for i := len(d.callbacks) - 1; i >= 0; i-- { 50 | d.callbacks[i]() 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /pkg/utils/xdefer/defers_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xdefer 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | // TestDefers test xdefer 28 | func TestDefers(t *testing.T) { 29 | defers := New() 30 | exitLog := "exit" 31 | fn1 := func() error { 32 | exitLog += " fn1" 33 | return nil 34 | } 35 | 36 | fn2 := func() error { 37 | exitLog += " fn2" 38 | return nil 39 | } 40 | 41 | fnN := func() error { 42 | exitLog += " fnN" 43 | return nil 44 | } 45 | 46 | defers.Add(fn1, fn2, fnN) 47 | defers.Close() 48 | assert.Equal(t, exitLog, "exit fnN fn2 fn1") 49 | } 50 | -------------------------------------------------------------------------------- /pkg/utils/xnet/ip.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xnet 20 | 21 | import ( 22 | "net" 23 | "strings" 24 | ) 25 | 26 | // ExternalIP return external ip 27 | func ExternalIP() (res []string) { 28 | inters, err := net.Interfaces() 29 | if err != nil { 30 | return 31 | } 32 | for _, inter := range inters { 33 | if !strings.HasPrefix(inter.Name, "lo") { 34 | addrs, err := inter.Addrs() 35 | if err != nil { 36 | continue 37 | } 38 | for _, addr := range addrs { 39 | if ipnet, ok := addr.(*net.IPNet); ok { 40 | if ipnet.IP.IsLoopback() || ipnet.IP.IsLinkLocalMulticast() || ipnet.IP.IsLinkLocalUnicast() { 41 | continue 42 | } 43 | if ip4 := ipnet.IP.To4(); ip4 != nil { 44 | switch true { 45 | case ip4[0] == 10: 46 | continue 47 | case ip4[0] == 172 && ip4[1] >= 16 && ip4[1] <= 31: 48 | continue 49 | case ip4[0] == 192 && ip4[1] == 168: 50 | continue 51 | default: 52 | res = append(res, ipnet.IP.String()) 53 | } 54 | } 55 | } 56 | } 57 | } 58 | } 59 | return 60 | } 61 | 62 | // InternalIP return internal ip 63 | func InternalIP() string { 64 | inters, err := net.Interfaces() 65 | if err != nil { 66 | return "" 67 | } 68 | for _, inter := range inters { 69 | if !isUp(inter.Flags) { 70 | continue 71 | } 72 | if !strings.HasPrefix(inter.Name, "lo") { 73 | addrs, err := inter.Addrs() 74 | if err != nil { 75 | continue 76 | } 77 | for _, addr := range addrs { 78 | if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() { 79 | if ipnet.IP.To4() != nil { 80 | return ipnet.IP.String() 81 | } 82 | } 83 | } 84 | } 85 | } 86 | return "" 87 | } 88 | 89 | // isUp Interface is up 90 | func isUp(v net.Flags) bool { 91 | return v&net.FlagUp == net.FlagUp 92 | } 93 | -------------------------------------------------------------------------------- /pkg/utils/xnet/ip_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xnet 20 | 21 | import ( 22 | "net" 23 | "testing" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | func TestInternalIP(t *testing.T) { 29 | ip := InternalIP() 30 | assert.True(t, len(ip) > 0) 31 | } 32 | 33 | func TestExternalIP(t *testing.T) { 34 | ip := ExternalIP() 35 | assert.True(t, len(ip) >= 0) 36 | } 37 | 38 | func TestIsUp(t *testing.T) { 39 | up := isUp(net.FlagUp) 40 | assert.Equal(t, true, up) 41 | } 42 | -------------------------------------------------------------------------------- /pkg/utils/xslice/float_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xslice 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestFloat64(t *testing.T) { 28 | f1 := []float64{3.12, 3.3, 3.0, 4, 7.222222} 29 | f2 := []float64{3.12, 3.3, 3.0, 4, 8.222222} 30 | 31 | assert.Equal(t, true, ContainFloat64(f1, 3.3)) 32 | assert.ElementsMatch(t, []float64{7.222222}, DiffFloat64(f1, f2)) 33 | assert.ElementsMatch(t, []float64{8.222222}, DiffFloat64(f2, f1)) 34 | assert.ElementsMatch(t, []float64{3.12, 3.3, 3.0, 4}, IntersectFloat64(f1, f2)) 35 | assert.ElementsMatch(t, []float64{3.12, 3.3, 3.0, 4, 7.222222, 3.12, 3.3, 3.0, 4, 8.222222}, MergeFloat64(f1, f2)) 36 | assert.ElementsMatch(t, []float64{3.12, 3.3, 4, 8.222222}, RemoveFloat64(f2, 3.0)) 37 | assert.ElementsMatch(t, []float64{8.222222, 4, 3.0, 3.3, 3.12}, ReverseFloat64(f2)) 38 | assert.ElementsMatch(t, []float64{3.0, 3.12, 3.3, 4, 7.222222}, SortFloat64(f1)) 39 | assert.ElementsMatch(t, []float64{3.0, 3.0, 3.12, 3.12, 3.3, 3.3, 4, 4, 7.222222, 8.222222}, SortFloat64(MergeFloat64(f1, f2))) 40 | assert.ElementsMatch(t, []float64{3.0, 3.12, 3.3, 4, 7.222222, 8.222222}, UniqueFloat64(MergeFloat64(f1, f2))) 41 | } 42 | -------------------------------------------------------------------------------- /pkg/utils/xslice/string.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xslice 20 | 21 | import ( 22 | "sort" 23 | ) 24 | 25 | // ContainString check if target is in ss 26 | func ContainString(s []string, target string) bool { 27 | for _, val := range s { 28 | if val == target { 29 | return true 30 | } 31 | } 32 | return false 33 | } 34 | 35 | // RemoveString remove empty target elements from ss 36 | func RemoveString(s []string, target string) []string { 37 | var ret = make([]string, 0) 38 | for _, val := range s { 39 | if val != target { 40 | ret = append(ret, val) 41 | } 42 | } 43 | return ret 44 | } 45 | 46 | // ReverseString reverse the input slice 47 | func ReverseString(s []string) []string { 48 | if len(s) < 2 { 49 | return s 50 | } 51 | 52 | start := 0 53 | end := len(s) - 1 54 | for start < end { 55 | tmp := s[start] 56 | s[start] = s[end] 57 | s[end] = tmp 58 | start++ 59 | end-- 60 | } 61 | return s 62 | } 63 | 64 | // DiffString computes the difference of two slices 65 | // return a slice containing all the entries from s1 but not in s2 66 | func DiffString(s1 []string, s2 []string) []string { 67 | ret := make([]string, 0) 68 | for _, val := range s1 { 69 | if !ContainString(s2, val) { 70 | ret = append(ret, val) 71 | } 72 | } 73 | return ret 74 | } 75 | 76 | // IntersectString computes the intersection of two slices 77 | // return a slice containing the entries exist in s1 and s2 78 | func IntersectString(s1 []string, s2 []string) []string { 79 | ret := make([]string, 0) 80 | for _, val := range s1 { 81 | if ContainString(s2, val) { 82 | ret = append(ret, val) 83 | } 84 | } 85 | return ret 86 | } 87 | 88 | // UniqueString Removes duplicate values from slice 89 | func UniqueString(s1 []string) []string { 90 | unique := make(map[string]struct{}) 91 | ret := make([]string, 0, len(s1)) 92 | for _, val := range s1 { 93 | if _, ok := unique[val]; ok { 94 | continue 95 | } 96 | unique[val] = struct{}{} 97 | ret = append(ret, val) 98 | } 99 | return ret 100 | } 101 | 102 | // MergeString merge one or more arrays 103 | func MergeString(s1 []string, s2 ...[]string) []string { 104 | if len(s2) == 0 { 105 | return s1 106 | } 107 | size := len(s1) 108 | for _, s := range s2 { 109 | size += len(s) 110 | } 111 | ret := make([]string, 0, size) 112 | ret = append(ret, s1...) 113 | for _, s := range s2 { 114 | ret = append(ret, s...) 115 | } 116 | return ret 117 | } 118 | 119 | // SortString sort string slice asc 120 | func SortString(s []string) []string { 121 | sort.Sort(sort.StringSlice(s)) 122 | return s 123 | } 124 | -------------------------------------------------------------------------------- /pkg/utils/xslice/string_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xslice 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestString(t *testing.T) { 28 | s1 := []string{"hello", "world", "hello+", "waterdrop"} 29 | s2 := []string{"hello", "world", "github", "waterdrop"} 30 | s3 := []string{"xstring"} 31 | 32 | assert.Equal(t, true, ContainString(s1, "hello")) 33 | assert.Equal(t, false, ContainString(s2, "hello+")) 34 | assert.ElementsMatch(t, []string{"hello+"}, DiffString(s1, s2)) 35 | assert.ElementsMatch(t, []string{"github"}, DiffString(s2, s1)) 36 | assert.ElementsMatch(t, []string{"hello", "world", "waterdrop"}, IntersectString(s1, s2)) 37 | assert.ElementsMatch(t, []string{"hello", "world", "hello+", "waterdrop", "hello", "world", "github", "waterdrop"}, MergeString(s1, s2)) 38 | assert.ElementsMatch(t, []string{"xstring"}, MergeString(s3)) 39 | assert.ElementsMatch(t, []string{"hello", "world", "hello+", "waterdrop", "hello", "world", "github", "waterdrop", "xstring"}, MergeString(s1, s2, s3)) 40 | assert.ElementsMatch(t, []string{"hello", "world", "hello+", "github", "waterdrop"}, UniqueString(MergeString(s1, s2))) 41 | assert.Equal(t, []string{"hello", "hello+", "waterdrop", "world"}, SortString(s1)) 42 | assert.ElementsMatch(t, []string{"hello", "world", "hello+"}, RemoveString(s1, "waterdrop")) 43 | assert.ElementsMatch(t, []string{"hello", "world", "hello+", "waterdrop"}, RemoveString(s1, "waterdrop+")) 44 | assert.ElementsMatch(t, []string{"waterdrop", "hello+", "world", "hello"}, ReverseString(s1)) 45 | } 46 | -------------------------------------------------------------------------------- /pkg/utils/xstring/string_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xstring 20 | 21 | import ( 22 | "bytes" 23 | "math/rand" 24 | "testing" 25 | 26 | "github.com/stretchr/testify/assert" 27 | ) 28 | 29 | func TestSubstr(t *testing.T) { 30 | assert.Equal(t, "waterdrop", Substr("Hello waterdrop!", 6, 15)) 31 | assert.Equal(t, "waterdrop!", Substr("Hello waterdrop!", 6, -1)) 32 | assert.Equal(t, "Hello waterdrop!", Substr("Hello waterdrop!", 0, -1)) 33 | assert.Equal(t, "Hello waterdrop!", Substr("Hello waterdrop!", 0, len("Hello waterdrop!")+1)) 34 | assert.NotEqual(t, "Hello waterdrop!", Substr("Hello waterdrop!", -1, len("Hello waterdrop!")+1)) 35 | assert.Equal(t, "", Substr("Hello waterdrop!", -1, len("Hello waterdrop!")+1)) 36 | assert.Equal(t, "", Substr("Hello waterdrop!", len("Hello waterdrop!")+1, len("Hello waterdrop!")+1)) 37 | } 38 | 39 | func TestReverse(t *testing.T) { 40 | assert.Equal(t, "!pordretaw olleH", Reverse("Hello waterdrop!")) 41 | assert.Equal(t, "World Peace", Reverse("ecaeP dlroW")) 42 | } 43 | 44 | func TestGenerateUUID(t *testing.T) { 45 | assert.Equal(t, 36, len(GenerateUUID())) 46 | } 47 | 48 | func TestStringToBytes(t *testing.T) { 49 | for i := 0; i < 100; i++ { 50 | s := RandomString(64) 51 | if !bytes.Equal([]byte(s), StringToBytes(s)) { 52 | assert.Fail(t, "don't match") 53 | } 54 | } 55 | } 56 | 57 | func TestBytesToString(t *testing.T) { 58 | data := make([]byte, 1024) 59 | for i := 0; i < 100; i++ { 60 | rand.Read(data) 61 | if string(data) != BytesToString(data) { 62 | assert.Fail(t, "don't match") 63 | } 64 | } 65 | } 66 | 67 | func TestStripContentType(t *testing.T) { 68 | assert.Equal(t, StripContentType("application/json;charset=utf-8"), "application/json") 69 | assert.Equal(t, StripContentType("application/json"), "application/json") 70 | } 71 | 72 | func TestGetLocaleLng(t *testing.T) { 73 | assert.Equal(t, _defalutLng, GetLocaleLng("*")) 74 | assert.Equal(t, _defalutLng, GetLocaleLng("")) 75 | assert.Equal(t, "fr", GetLocaleLng("fr-CH, fr;q=0.9")) 76 | assert.Equal(t, "en", GetLocaleLng("en;q=0.8")) 77 | assert.Equal(t, _defalutLng, GetLocaleLng("*;q=0.5")) 78 | assert.Equal(t, "zh", GetLocaleLng(" zh-CN,zh;q=0.5")) 79 | } 80 | 81 | // BenchmarkRandomString bench generate random length string 82 | func BenchmarkRandomString(b *testing.B) { 83 | b.RunParallel(func(pb *testing.PB) { 84 | for pb.Next() { 85 | _ = RandomString(16) 86 | } 87 | }) 88 | } 89 | -------------------------------------------------------------------------------- /pkg/utils/xsync/singlecall.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xsync 20 | 21 | import "sync" 22 | 23 | // inflight is an in-flight or completed Do call 24 | type inflight struct { 25 | wg sync.WaitGroup 26 | val interface{} 27 | err error 28 | } 29 | 30 | // Group represents a class of work and forms a namespace in which 31 | // units of work can be executed with duplicate suppression. 32 | type SingleCall struct { 33 | mu sync.Mutex // protects m 34 | m map[string]*inflight // lazily initialized 35 | } 36 | 37 | // Do executes and returns the results of the given function, making 38 | // sure that only one execution is in-flight for a given key at a 39 | // time. If a duplicate comes in, the duplicate caller waits for the 40 | // original to complete and receives the same results. 41 | func (s *SingleCall) Do(key string, fn func() (interface{}, error)) (interface{}, error) { 42 | s.mu.Lock() 43 | if s.m == nil { 44 | s.m = make(map[string]*inflight) 45 | } 46 | if c, ok := s.m[key]; ok { 47 | s.mu.Unlock() 48 | c.wg.Wait() 49 | return c.val, c.err 50 | } 51 | c := &inflight{} 52 | c.wg.Add(1) 53 | s.m[key] = c 54 | s.mu.Unlock() 55 | 56 | c.val, c.err = fn() 57 | c.wg.Done() 58 | 59 | s.mu.Lock() 60 | delete(s.m, key) 61 | s.mu.Unlock() 62 | 63 | return c.val, c.err 64 | } 65 | -------------------------------------------------------------------------------- /pkg/utils/xsync/singlecall_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xsync 20 | 21 | import ( 22 | "errors" 23 | "sync/atomic" 24 | "testing" 25 | 26 | "github.com/stretchr/testify/assert" 27 | ) 28 | 29 | func TestDo(t *testing.T) { 30 | sc := &SingleCall{} 31 | v, err := sc.Do("key", func() (interface{}, error) { 32 | return "inflight", nil 33 | }) 34 | 35 | assert.Equal(t, err, nil) 36 | assert.Equal(t, v.(string), "inflight") 37 | } 38 | 39 | func TestDoErr(t *testing.T) { 40 | sc := &SingleCall{} 41 | someErr := errors.New("some error") 42 | v, err := sc.Do("key", func() (interface{}, error) { 43 | return nil, someErr 44 | }) 45 | assert.Equal(t, nil, v) 46 | assert.Equal(t, err.Error(), "some error") 47 | } 48 | 49 | func TestDoDupSuppress(t *testing.T) { 50 | sc := &SingleCall{} 51 | var calls int32 52 | fn := func() (interface{}, error) { 53 | atomic.AddInt32(&calls, 1) 54 | return "inflight", nil 55 | } 56 | 57 | wg := &WaitGroupWrapper{} 58 | for i := 0; i < 10; i++ { 59 | wg.Wrap(func() { 60 | v, err := sc.Do("key", fn) 61 | if err != nil { 62 | t.Errorf("Do error: %v", err) 63 | } 64 | if v.(string) != "inflight" { 65 | t.Errorf("got %q; want %q", v, "bar") 66 | } 67 | }) 68 | } 69 | wg.Wait() 70 | assert.NotEqual(t, 10, calls) 71 | } 72 | -------------------------------------------------------------------------------- /pkg/utils/xsync/waitgroup.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xsync 20 | 21 | import "sync" 22 | 23 | // WaitGroupWrapper wrap wait group 24 | type WaitGroupWrapper struct { 25 | sync.WaitGroup 26 | } 27 | 28 | // Wrap wrap func to wait group 29 | func (wg *WaitGroupWrapper) Wrap(cb func()) { 30 | wg.Add(1) 31 | go func() { 32 | cb() 33 | wg.Done() 34 | }() 35 | } 36 | -------------------------------------------------------------------------------- /pkg/utils/xsync/waitgroup_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xsync 20 | 21 | import ( 22 | "sync/atomic" 23 | "testing" 24 | 25 | "github.com/stretchr/testify/assert" 26 | ) 27 | 28 | func TestWrap(t *testing.T) { 29 | wg := &WaitGroupWrapper{} 30 | 31 | var result int32 32 | for i := 0; i <= 100; i++ { 33 | j := i 34 | wg.Wrap(func() { 35 | atomic.AddInt32(&result, int32(j)) 36 | }) 37 | } 38 | wg.Wait() 39 | assert.Equal(t, int32(5050), result) 40 | } 41 | -------------------------------------------------------------------------------- /pkg/utils/xtime/jitter.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2023 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xtime 20 | 21 | import ( 22 | "math" 23 | "math/rand" 24 | "time" 25 | ) 26 | 27 | const coefficient = 0.1 28 | 29 | // JitterTime simulates jitter by scaling a time.Duration randomly within factor 30 | // The duration d must be greater than zero; and the scaling factor must be 31 | // within the range 0 < factor <= 1.0, default factor is 0.1 32 | func JitterTime(d time.Duration, factor ...float64) (scale time.Duration) { 33 | if d <= 0 { 34 | return d 35 | } 36 | 37 | f := coefficient 38 | if len(factor) > 0 { 39 | if factor[0] > 1.0 || factor[0] <= 0 { 40 | return d 41 | } 42 | f = factor[0] 43 | } 44 | 45 | min, max := bounds(int64(d), f) 46 | scale = time.Duration(randRange(min, max)) 47 | return 48 | } 49 | 50 | // Jitter simulates jitter by scaling a integer randomly within factor 51 | // The duration d must be greater than zero; and the scaling factor must be 52 | // within the range 0 < factor <= 1.0, default factor is 0.1 53 | func Jitter(n int64, factor ...float64) (scale int64) { 54 | if n <= 0 { 55 | return n 56 | } 57 | 58 | f := coefficient 59 | if len(factor) > 0 { 60 | if factor[0] > 1.0 || factor[0] <= 0 { 61 | return n 62 | } 63 | f = factor[0] 64 | } 65 | 66 | min, max := bounds(n, f) 67 | scale = randRange(min, max) 68 | return 69 | } 70 | 71 | // bounds returns the min and max values for n after applying scaling factor 72 | // if the max overflow, then truncate and return math.MaxInt64 73 | func bounds(n int64, factor ...float64) (min, max int64) { 74 | f := coefficient 75 | if len(factor) > 0 { 76 | f = factor[0] 77 | } 78 | 79 | minf := math.Floor(float64(n) * (1 - f)) 80 | maxf := math.Ceil(float64(n) * (1 + f)) 81 | 82 | if maxf > math.MaxInt64 { 83 | return int64(minf), math.MaxInt64 84 | } 85 | return int64(minf), int64(maxf) 86 | } 87 | 88 | // randRange returns a non-negative pseudo-random number in the half open 89 | // interval [min, max) from the default source 90 | func randRange(min, max int64) int64 { 91 | if min == max { 92 | return min 93 | } 94 | return rand.Int63n(max-min) + min 95 | } 96 | -------------------------------------------------------------------------------- /pkg/utils/xtime/jitter_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2023 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package xtime 20 | 21 | import ( 22 | "testing" 23 | "time" 24 | ) 25 | 26 | func TestJitterTime(t *testing.T) { 27 | t.Run("sample test", func(t *testing.T) { 28 | const ( 29 | d = 10 * time.Second 30 | f = 0.5 31 | samples = 20 32 | ) 33 | 34 | for i := 0; i < samples; i++ { 35 | r := JitterTime(d, f) 36 | t.Log(r) 37 | if r < 5*time.Second || r > 15*time.Second { 38 | t.Error("sample outside of range: ", r) 39 | } 40 | } 41 | 42 | for i := 0; i < samples; i++ { 43 | r := JitterTime(d) 44 | t.Log(r) 45 | if r < 9*time.Second || r > 11*time.Second { 46 | t.Error("sample outside of range: ", r) 47 | } 48 | } 49 | }) 50 | } 51 | 52 | func TestJitter(t *testing.T) { 53 | t.Run("sample test", func(t *testing.T) { 54 | const ( 55 | d = 10 56 | f = 0.5 57 | samples = 20 58 | ) 59 | 60 | for i := 0; i < samples; i++ { 61 | r := Jitter(d, f) 62 | t.Log(r) 63 | if r < 5 || r > 15 { 64 | t.Error("sample outside of range: ", r) 65 | } 66 | } 67 | 68 | for i := 0; i < samples; i++ { 69 | r := Jitter(d) 70 | t.Log(r) 71 | if r < 9 || r > 11 { 72 | t.Error("sample outside of range: ", r) 73 | } 74 | } 75 | }) 76 | } 77 | -------------------------------------------------------------------------------- /tests/proto/demo/demo.proto: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | // 定义项目 API 的 proto 文件 可以同时描述 gRPC 和 HTTP API 20 | // protobuf 文件参考: 21 | // - https://developers.google.com/protocol-buffers/ 22 | syntax = "proto3"; 23 | 24 | import "google/protobuf/empty.proto"; 25 | 26 | // package 命名使用 {appid}.{version} 的方式, version 形如 v1, v2 .. 27 | package service.demo.v1; 28 | 29 | // NOTE: 最后请删除这些无用的注释 (゜-゜)つロ 30 | 31 | option go_package = "./;demo"; 32 | 33 | service Demo { 34 | rpc SayHello(HelloReq) returns (.google.protobuf.Empty); 35 | rpc SayHelloURL(HelloReq) returns (HelloResp); 36 | } 37 | 38 | message HelloReq { 39 | string name = 1; 40 | } 41 | 42 | message HelloResp { 43 | string Content = 1; 44 | } 45 | -------------------------------------------------------------------------------- /tools/waterdrop/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/UnderTreeTech/waterdrop/tools/waterdrop 2 | 3 | go 1.14 4 | 5 | require ( 6 | github.com/fatih/color v1.12.0 7 | github.com/urfave/cli/v2 v2.3.0 8 | golang.org/x/mod v0.5.0 9 | golang.org/x/tools v0.1.5 10 | ) 11 | -------------------------------------------------------------------------------- /tools/waterdrop/main.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package main 20 | 21 | import ( 22 | "fmt" 23 | "log" 24 | "os" 25 | 26 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/ecode" 27 | 28 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/upgrade" 29 | 30 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/project" 31 | 32 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/swagger" 33 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/testgen/utgen" 34 | 35 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/protoc" 36 | 37 | "github.com/urfave/cli/v2" 38 | ) 39 | 40 | const Version = "v0.2.0" 41 | 42 | // main tool entry point 43 | func main() { 44 | app := cli.NewApp() 45 | app.Name = "waterdrop" 46 | app.Usage = "waterdrop tools" 47 | app.Version = Version 48 | app.Commands = []*cli.Command{ 49 | project.ProjectCmd, 50 | protoc.ProtocCmd, 51 | swagger.SwaggerCmd, 52 | utgen.UTCmd, 53 | upgrade.UpgradeCmd, 54 | ecode.EcodeCmd, 55 | } 56 | 57 | if err := app.Run(os.Args); err != nil { 58 | log.Fatalf(fmt.Sprintf("run waterdrop tool fail, error is %s", err.Error())) 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /tools/waterdrop/project/README.md: -------------------------------------------------------------------------------- 1 | ## New Project Layout 2 | 3 | This idea is borrowed from [kratos](https://github.com/go-kratos/kratos) but make some minor changes. 4 | Thank you for kratos contributors. -------------------------------------------------------------------------------- /tools/waterdrop/project/base/install.go: -------------------------------------------------------------------------------- 1 | //go:build go1.17 2 | // +build go1.17 3 | 4 | package base 5 | 6 | import ( 7 | "fmt" 8 | "os" 9 | "os/exec" 10 | ) 11 | 12 | // GoInstall go get path. 13 | func GoInstall(path ...string) error { 14 | for _, p := range path { 15 | fmt.Printf("go install %s@latest\n", p) 16 | cmd := exec.Command("go", "install", fmt.Sprintf("%s@latest", p)) 17 | cmd.Stdout = os.Stdout 18 | cmd.Stderr = os.Stderr 19 | if err := cmd.Run(); err != nil { 20 | return err 21 | } 22 | } 23 | return nil 24 | } 25 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/install_compatible.go: -------------------------------------------------------------------------------- 1 | //go:build !go1.17 2 | // +build !go1.17 3 | 4 | package base 5 | 6 | import ( 7 | "fmt" 8 | "os" 9 | "os/exec" 10 | ) 11 | 12 | // GoInstall go get path. 13 | func GoInstall(path ...string) error { 14 | for _, p := range path { 15 | fmt.Printf("go get -u %s\n", p) 16 | cmd := exec.Command("go", "get", "-u", p) 17 | cmd.Stdout = os.Stdout 18 | cmd.Stderr = os.Stderr 19 | if err := cmd.Run(); err != nil { 20 | return err 21 | } 22 | } 23 | return nil 24 | } 25 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/mod.go: -------------------------------------------------------------------------------- 1 | package base 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "io/ioutil" 7 | "os/exec" 8 | "strings" 9 | 10 | "golang.org/x/mod/modfile" 11 | ) 12 | 13 | // ModulePath returns go module path. 14 | func ModulePath(filename string) (string, error) { 15 | modBytes, err := ioutil.ReadFile(filename) 16 | if err != nil { 17 | return "", err 18 | } 19 | return modfile.ModulePath(modBytes), nil 20 | } 21 | 22 | // ModuleVersion returns module version. 23 | func ModuleVersion(path string) (string, error) { 24 | stdout := &bytes.Buffer{} 25 | fd := exec.Command("go", "mod", "graph") 26 | fd.Stdout = stdout 27 | fd.Stderr = stdout 28 | if err := fd.Run(); err != nil { 29 | return "", err 30 | } 31 | rd := bufio.NewReader(stdout) 32 | for { 33 | line, _, err := rd.ReadLine() 34 | if err != nil { 35 | return "", err 36 | } 37 | str := string(line) 38 | i := strings.Index(str, "@") 39 | if strings.Contains(str, path+"@") && i != -1 { 40 | return path + str[i:], nil 41 | } 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/mod_test.go: -------------------------------------------------------------------------------- 1 | package base 2 | 3 | import "testing" 4 | 5 | func TestModuleVersion(t *testing.T) { 6 | v, err := ModuleVersion("golang.org/x/mod") 7 | if err != nil { 8 | t.Fatal(err) 9 | } 10 | t.Log(v) 11 | } 12 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/path.go: -------------------------------------------------------------------------------- 1 | package base 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io/ioutil" 7 | "log" 8 | "os" 9 | "path" 10 | "path/filepath" 11 | "strings" 12 | 13 | "github.com/fatih/color" 14 | ) 15 | 16 | func waterdropHome() string { 17 | dir, err := os.UserHomeDir() 18 | if err != nil { 19 | log.Fatal(err) 20 | } 21 | home := path.Join(dir, ".waterdrop") 22 | if _, err := os.Stat(home); os.IsNotExist(err) { 23 | if err := os.MkdirAll(home, 0700); err != nil { 24 | log.Fatal(err) 25 | } 26 | } 27 | return home 28 | } 29 | 30 | func waterdropHomeWithDir(dir string) string { 31 | home := path.Join(waterdropHome(), dir) 32 | if _, err := os.Stat(home); os.IsNotExist(err) { 33 | if err := os.MkdirAll(home, 0700); err != nil { 34 | log.Fatal(err) 35 | } 36 | } 37 | return home 38 | } 39 | 40 | func copyFile(src, dst string, replaces []string) error { 41 | var err error 42 | srcinfo, err := os.Stat(src) 43 | if err != nil { 44 | return err 45 | } 46 | buf, err := ioutil.ReadFile(src) 47 | if err != nil { 48 | return err 49 | } 50 | var old string 51 | for i, next := range replaces { 52 | if i%2 == 0 { 53 | old = next 54 | continue 55 | } 56 | buf = bytes.ReplaceAll(buf, []byte(old), []byte(next)) 57 | } 58 | return ioutil.WriteFile(dst, buf, srcinfo.Mode()) 59 | } 60 | 61 | func copyDir(src, dst string, replaces, ignores []string) error { 62 | var err error 63 | var fds []os.FileInfo 64 | var srcinfo os.FileInfo 65 | 66 | if srcinfo, err = os.Stat(src); err != nil { 67 | return err 68 | } 69 | 70 | if err = os.MkdirAll(dst, srcinfo.Mode()); err != nil { 71 | return err 72 | } 73 | 74 | if fds, err = ioutil.ReadDir(src); err != nil { 75 | return err 76 | } 77 | for _, fd := range fds { 78 | if hasSets(fd.Name(), ignores) { 79 | continue 80 | } 81 | srcfp := path.Join(src, fd.Name()) 82 | dstfp := path.Join(dst, fd.Name()) 83 | if fd.IsDir() { 84 | if err = copyDir(srcfp, dstfp, replaces, ignores); err != nil { 85 | return err 86 | } 87 | } else { 88 | if err = copyFile(srcfp, dstfp, replaces); err != nil { 89 | return err 90 | } 91 | } 92 | } 93 | return nil 94 | } 95 | 96 | func hasSets(name string, sets []string) bool { 97 | for _, ig := range sets { 98 | if ig == name { 99 | return true 100 | } 101 | } 102 | return false 103 | } 104 | 105 | func Tree(path string, dir string) { 106 | filepath.Walk(path, func(path string, info os.FileInfo, err error) error { 107 | if !info.IsDir() { 108 | fmt.Printf("%s %s (%v bytes)\n", color.GreenString("CREATED"), strings.Replace(path, dir+"/", "", -1), info.Size()) 109 | } 110 | return nil 111 | }) 112 | } 113 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/repo.go: -------------------------------------------------------------------------------- 1 | package base 2 | 3 | import ( 4 | "context" 5 | "os" 6 | "os/exec" 7 | "path" 8 | "strings" 9 | ) 10 | 11 | // Repo is git repository manager. 12 | type Repo struct { 13 | url string 14 | home string 15 | branch string 16 | } 17 | 18 | // NewRepo new a repository manager. 19 | func NewRepo(url string, branch string) *Repo { 20 | var start int 21 | start = strings.Index(url, "//") 22 | if start == -1 { 23 | start = strings.Index(url, ":") + 1 24 | } else { 25 | start += 2 26 | } 27 | end := strings.LastIndex(url, "/") 28 | return &Repo{ 29 | url: url, 30 | home: waterdropHomeWithDir("repo/" + url[start:end]), 31 | branch: branch, 32 | } 33 | } 34 | 35 | // Path returns the repository cache path. 36 | func (r *Repo) Path() string { 37 | start := strings.LastIndex(r.url, "/") 38 | end := strings.LastIndex(r.url, ".git") 39 | if end == -1 { 40 | end = len(r.url) 41 | } 42 | branch := "" 43 | if r.branch == "" { 44 | branch = "@main" 45 | } else { 46 | branch = "@" + r.branch 47 | } 48 | return path.Join(r.home, r.url[start+1:end]+branch) 49 | } 50 | 51 | // Pull fetch the repository from remote url. 52 | func (r *Repo) Pull(ctx context.Context) error { 53 | cmd := exec.Command("git", "symbolic-ref", "HEAD") 54 | cmd.Dir = r.Path() 55 | err := cmd.Run() 56 | if err != nil { 57 | return nil 58 | } 59 | cmd = exec.Command("git", "pull") 60 | cmd.Dir = r.Path() 61 | cmd.Stderr = os.Stderr 62 | cmd.Stdout = os.Stdout 63 | err = cmd.Run() 64 | return err 65 | } 66 | 67 | // Clone clones the repository to cache path. 68 | func (r *Repo) Clone(ctx context.Context) error { 69 | if _, err := os.Stat(r.Path()); !os.IsNotExist(err) { 70 | return r.Pull(ctx) 71 | } 72 | cmd := &exec.Cmd{} 73 | if r.branch == "" { 74 | cmd = exec.Command("git", "clone", r.url, r.Path()) 75 | } else { 76 | cmd = exec.Command("git", "clone", "-b", r.branch, r.url, r.Path()) 77 | } 78 | cmd.Stderr = os.Stderr 79 | cmd.Stdout = os.Stdout 80 | err := cmd.Run() 81 | return err 82 | } 83 | 84 | // CopyTo copies the repository to project path. 85 | func (r *Repo) CopyTo(ctx context.Context, to string, modPath string, ignores []string) error { 86 | if err := r.Clone(ctx); err != nil { 87 | return err 88 | } 89 | mod, err := ModulePath(path.Join(r.Path(), "go.mod")) 90 | if err != nil { 91 | return err 92 | } 93 | return copyDir(r.Path(), to, []string{mod, modPath}, ignores) 94 | } 95 | -------------------------------------------------------------------------------- /tools/waterdrop/project/base/repo_test.go: -------------------------------------------------------------------------------- 1 | package base 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | ) 7 | 8 | func TestRepo(t *testing.T) { 9 | r := NewRepo("https://github.com/go-kratos/service-layout.git", "") 10 | if err := r.Clone(context.Background()); err != nil { 11 | t.Fatal(err) 12 | } 13 | if err := r.CopyTo(context.Background(), "/tmp/test_repo", "github.com/go-kratos/kratos-layout", nil); err != nil { 14 | t.Fatal(err) 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /tools/waterdrop/project/layout.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package project 20 | 21 | import ( 22 | "context" 23 | "fmt" 24 | "os" 25 | "path" 26 | "time" 27 | 28 | "github.com/fatih/color" 29 | 30 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/project/base" 31 | 32 | "github.com/urfave/cli/v2" 33 | ) 34 | 35 | var ProjectCmd = &cli.Command{ 36 | Name: "new", 37 | Usage: "waterdrop project layout tools", 38 | Action: run, 39 | SkipFlagParsing: false, 40 | UsageText: "new", 41 | } 42 | 43 | func run(c *cli.Context) (err error) { 44 | pwd, err := os.Getwd() 45 | if err != nil { 46 | panic(err) 47 | } 48 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute) 49 | defer cancel() 50 | 51 | if c.Args().Len() == 0 { 52 | fmt.Println("You must assign project name") 53 | return 54 | } 55 | 56 | projectName := c.Args().Slice()[0] 57 | to := path.Join(pwd, projectName) 58 | if _, err = os.Stat(to); !os.IsNotExist(err) { 59 | fmt.Println(fmt.Sprintf("Project %s already exists, please reassign a new name.", color.RedString(projectName))) 60 | return nil 61 | } 62 | 63 | repo := base.NewRepo("https://github.com/UnderTreeTech/layout.git", "master") 64 | if err = repo.CopyTo(ctx, to, path.Base(projectName), []string{".git", ".github"}); err != nil { 65 | fmt.Println("New project fail, please try again later.") 66 | return 67 | } 68 | base.Tree(to, pwd) 69 | 70 | fmt.Println("\n🍺 🍺 🍺 Project creation succeeded", color.RedString(projectName), "🍺 🍺 🍺 ") 71 | fmt.Print("💻 Use the following command to start the project 👇:\n\n") 72 | fmt.Println(color.WhiteString("$ cd %s", projectName)) 73 | fmt.Println(color.WhiteString("$ go mod tidy")) 74 | fmt.Println(color.WhiteString("$ cd cmd")) 75 | fmt.Println(color.WhiteString("$ go build -o %s main.go", projectName)) 76 | fmt.Println(color.WhiteString("$ ./%s -conf=../configs/application.toml\n", projectName)) 77 | fmt.Println("🤝 🤝 🤝 Thanks for using Waterdrop 🤝 🤝 🤝 ") 78 | return 79 | } 80 | -------------------------------------------------------------------------------- /tools/waterdrop/protoc/action.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package protoc 20 | 21 | import ( 22 | "github.com/urfave/cli/v2" 23 | ) 24 | 25 | // run run commands 26 | func run(ctx *cli.Context) (err error) { 27 | if err = checkProtocEnv(); err != nil { 28 | return 29 | } 30 | 31 | // 根据指定目录下的proto 文件 生成pb.go 文件 32 | if genGRPC { 33 | if err = generateGRPC(ctx); err != nil { 34 | return 35 | } 36 | } 37 | 38 | // 根据指定目录下的proto 文件 生成pb.swagger.json 文件 39 | if genSwagger { 40 | if err = generateSwagger(ctx); err != nil { 41 | return 42 | } 43 | } 44 | 45 | return 46 | } 47 | -------------------------------------------------------------------------------- /tools/waterdrop/protoc/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package protoc 20 | 21 | import "github.com/urfave/cli/v2" 22 | 23 | const ( 24 | protocHelpTemplate = `proto tools` 25 | ) 26 | 27 | var ( 28 | genGRPC bool 29 | genSwagger bool 30 | ) 31 | 32 | var ProtocCmd = &cli.Command{ 33 | Name: "protoc", 34 | Usage: "waterdrop protoc tools", 35 | Action: run, 36 | SkipFlagParsing: false, 37 | UsageText: protocHelpTemplate, 38 | Flags: []cli.Flag{ 39 | &cli.BoolFlag{ 40 | Name: "grpc", 41 | Usage: "whether to generate GRPC code", 42 | Destination: &genGRPC, 43 | }, 44 | &cli.BoolFlag{ 45 | Name: "swagger", 46 | Usage: "whether to use swagger for generation", 47 | Destination: &genSwagger, 48 | }, 49 | }, 50 | } 51 | -------------------------------------------------------------------------------- /tools/waterdrop/protoc/common.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package protoc 20 | 21 | import ( 22 | "errors" 23 | "fmt" 24 | "go/build" 25 | "os" 26 | "os/exec" 27 | "path/filepath" 28 | "runtime" 29 | "strings" 30 | 31 | "github.com/urfave/cli/v2" 32 | ) 33 | 34 | // download protoc: https://github.com/protocolbuffers/protobuf/releases. 35 | // copy protoc and include directory to GOPATH/bin 36 | // install protoc-gen-go, go get -u github.com/golang/protobuf/protoc-gen-go 37 | // checkProtocEnv check if exist protoc 38 | func checkProtocEnv() (err error) { 39 | if _, err = exec.LookPath("protoc"); err != nil { 40 | err = errors.New("You haven't installed Protobuf yet，Please visit this page to install with your own system：https://github.com/protocolbuffers/protobuf/releases") 41 | return err 42 | } 43 | return nil 44 | } 45 | 46 | // doGenerate do generate file command 47 | func doGenerate(ctx *cli.Context, protocCmd string) (err error) { 48 | files := ctx.Args().Slice() 49 | if len(files) == 0 { 50 | files, _ = filepath.Glob("*.proto") 51 | } 52 | 53 | pwd, _ := os.Getwd() 54 | // case go path 55 | var contectflag string 56 | // gopath setting could be many params, such as 57 | // { 58 | // win : "C:\go\src;D:\go\src" 59 | // linux: "/home/go:/root/go" 60 | // } 61 | if runtime.GOOS == "windows" { 62 | contectflag = ";" 63 | } else { 64 | contectflag = ":" 65 | } 66 | gosrcarr := strings.Split(build.Default.GOPATH, contectflag) 67 | if len(gosrcarr) < 1 { 68 | fmt.Println("gopath directory does not exist, please create it in your GOPATH") 69 | return nil 70 | } 71 | gosrc := filepath.Join(gosrcarr[0], "src") 72 | _, err = os.Stat(gosrc) 73 | if err != nil { 74 | fmt.Printf("src directory does not exist, please create it in your GOPATH: %v", gosrcarr[0]) 75 | return nil 76 | } 77 | 78 | cmdLine := fmt.Sprintf(protocCmd, gosrc, pwd) 79 | args := strings.Split(cmdLine, " ") 80 | args = append(args, files...) 81 | cmd := exec.Command(args[0], args[1:]...) 82 | cmd.Dir = pwd 83 | cmd.Env = os.Environ() 84 | cmd.Stdout = os.Stdout 85 | cmd.Stderr = os.Stderr 86 | err = cmd.Run() 87 | 88 | return 89 | } 90 | -------------------------------------------------------------------------------- /tools/waterdrop/protoc/grpc.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package protoc 20 | 21 | import ( 22 | "github.com/urfave/cli/v2" 23 | ) 24 | 25 | const ( 26 | //default generate pb files to current directory 27 | _grpcProtocCmd = `protoc -I=%s -I=%s --go_out=plugins=grpc:.` 28 | ) 29 | 30 | // generateGRPC generate grpc pb files 31 | func generateGRPC(ctx *cli.Context) error { 32 | if err := doGenerate(ctx, _grpcProtocCmd); err != nil { 33 | return err 34 | } 35 | 36 | return nil 37 | } 38 | -------------------------------------------------------------------------------- /tools/waterdrop/protoc/swagger.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package protoc 20 | 21 | import ( 22 | "os/exec" 23 | 24 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/utils" 25 | 26 | "github.com/urfave/cli/v2" 27 | ) 28 | 29 | const ( 30 | // go get protoc-gen-swagger 31 | _getSwaggerGen = "go get -u github.com/grpc-ecosystem/grpc-gateway/protoc-gen-swagger" 32 | // default generate swagger json files to current directory 33 | _swaggerProtoc = `protoc -I=%s -I=%s --swagger_out=logtostderr=true,generate_unbound_methods=true,disable_default_errors=true:.` 34 | ) 35 | 36 | // installSwaggerProtoc install protoc-gen-swagger 37 | func installSwaggerProtoc() error { 38 | if _, err := exec.LookPath("protoc-gen-swagger"); err != nil { 39 | if err := utils.ExecuteGoGet(_getSwaggerGen); err != nil { 40 | return err 41 | } 42 | } 43 | return nil 44 | } 45 | 46 | // generateSwagger generate swagger json file 47 | func generateSwagger(ctx *cli.Context) error { 48 | if err := installSwaggerProtoc(); err != nil { 49 | return err 50 | } 51 | 52 | if err := doGenerate(ctx, _swaggerProtoc); err != nil { 53 | return err 54 | } 55 | 56 | return nil 57 | } 58 | -------------------------------------------------------------------------------- /tools/waterdrop/swagger/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package swagger 20 | 21 | import ( 22 | "os" 23 | "os/exec" 24 | 25 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/utils" 26 | 27 | "github.com/urfave/cli/v2" 28 | ) 29 | 30 | var _installSwagger = `go get -u github.com/go-swagger/go-swagger/cmd/swagger` 31 | 32 | var SwaggerCmd = &cli.Command{ 33 | Name: "swagger", 34 | Usage: "waterdrop swagger tools", 35 | Action: run, 36 | SkipFlagParsing: false, 37 | UsageText: "swagger", 38 | } 39 | 40 | // run execute swagger serve command 41 | func run(ctx *cli.Context) error { 42 | if _, err := exec.LookPath("swagger"); err != nil { 43 | if err = utils.ExecuteGoGet(_installSwagger); err != nil { 44 | return err 45 | } 46 | } 47 | 48 | pwd, _ := os.Getwd() 49 | return utils.RunTool(ctx, pwd, ctx.Args().Slice()) 50 | } 51 | -------------------------------------------------------------------------------- /tools/waterdrop/testgen/common/utils.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package common 20 | 21 | import ( 22 | "fmt" 23 | "strings" 24 | 25 | "golang.org/x/tools/imports" 26 | ) 27 | 28 | const GenTestSuccess = ` 29 | 江城子·程序记梦 30 | 十年编程两茫茫，工期短，需求长。 31 | 百行代码，Bug到处藏。 32 | 纵使上线又如何，新版本，继续忙。 33 | 黑白颠倒没商量，趴桌上，进梦乡。 34 | 夜半梦醒，无人在身旁。 35 | 最怕灯火阑珊时，手机响，心里慌。 36 | // Generation success. Powered by Waterdrop 37 | ` 38 | 39 | // GoImport Use golang.org/x/tools/imports auto import pkg 40 | func GoImport(file string, bytes []byte) (res []byte, err error) { 41 | options := &imports.Options{ 42 | TabWidth: 8, 43 | TabIndent: true, 44 | Comments: true, 45 | Fragment: true, 46 | } 47 | if res, err = imports.Process(file, bytes, options); err != nil { 48 | fmt.Printf("GoImport(%s) error(%v)", file, err) 49 | res = bytes 50 | return 51 | } 52 | return 53 | } 54 | 55 | // ConvertMethod checkout the file belongs to dao or not 56 | func ConvertMethod(path string) (method string) { 57 | switch { 58 | case strings.Contains(path, "/dao"): 59 | method = "d" 60 | case strings.Contains(path, "/service"): 61 | method = "s" 62 | default: 63 | method = "" 64 | } 65 | return 66 | } 67 | -------------------------------------------------------------------------------- /tools/waterdrop/testgen/utgen/cmd.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package utgen 20 | 21 | import ( 22 | "fmt" 23 | 24 | "github.com/fatih/color" 25 | 26 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/testgen/common" 27 | 28 | "github.com/urfave/cli/v2" 29 | ) 30 | 31 | var genFunc string 32 | 33 | var UTCmd = &cli.Command{ 34 | Name: "utgen", 35 | Usage: "waterdrop unit test tools", 36 | Action: run, 37 | SkipFlagParsing: false, 38 | UsageText: "ut", 39 | Flags: []cli.Flag{ 40 | &cli.StringFlag{ 41 | Name: "func", 42 | Usage: "whether to generate unit test by func", 43 | Destination: &genFunc, 44 | }, 45 | }, 46 | } 47 | 48 | // run run generate unit test command 49 | func run(ctx *cli.Context) error { 50 | var ( 51 | err error 52 | files []string 53 | parses []*common.Parse 54 | ) 55 | 56 | if err = common.ParseArgs(ctx.Args().Slice(), &files, 0); err != nil { 57 | panic(err) 58 | } 59 | 60 | if parses, err = common.ParseFile(files...); err != nil { 61 | panic(err) 62 | } 63 | 64 | if err = genTest(parses); err != nil { 65 | panic(err) 66 | } 67 | 68 | fmt.Println(color.GreenString(common.GenTestSuccess)) 69 | return nil 70 | } 71 | -------------------------------------------------------------------------------- /tools/waterdrop/testgen/utgen/templete.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package utgen 20 | 21 | var ( 22 | tpPackage = "package %s\n\n" 23 | tpImport = "import (\n\t%s\n)\n\n" 24 | tpVar = "var (\n\t%s\n)\n" 25 | tpTestReset = "\n\t\tReset(func() {%s\n\t\t})" 26 | tpTestFunc = "\nfunc Test%s%s(t *testing.T){%s\n\tConvey(\"%s\", t, func(){\n\t\t%s\tConvey(\"When everything goes positive\", func(){\n\t\t\t%s\n\t\t\t})\n\t\t})%s\n\t})\n}\n\n" 27 | ) 28 | -------------------------------------------------------------------------------- /tools/waterdrop/upgrade/upgrade.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2021 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package upgrade 20 | 21 | import ( 22 | "fmt" 23 | 24 | "github.com/fatih/color" 25 | 26 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop/project/base" 27 | "github.com/urfave/cli/v2" 28 | ) 29 | 30 | var UpgradeCmd = &cli.Command{ 31 | Name: "upgrade", 32 | Usage: "waterdrop upgrade tools", 33 | Action: run, 34 | SkipFlagParsing: false, 35 | UsageText: "upgrade", 36 | } 37 | 38 | func run(c *cli.Context) (err error) { 39 | err = base.GoInstall( 40 | "github.com/UnderTreeTech/waterdrop/tools/waterdrop", 41 | ) 42 | 43 | if err != nil { 44 | fmt.Println("upgrade waterdrop fail", err.Error()) 45 | return 46 | } 47 | 48 | fmt.Println(color.GreenString("upgrade tool waterdrop success")) 49 | return 50 | } 51 | -------------------------------------------------------------------------------- /tools/waterdrop/utils/env.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package utils 20 | 21 | import ( 22 | "bytes" 23 | "fmt" 24 | "io/ioutil" 25 | "os" 26 | "path/filepath" 27 | "sync" 28 | ) 29 | 30 | var envCache struct { 31 | once sync.Once 32 | m map[string]string 33 | } 34 | 35 | // EnvFile returns the name of the Go environment configuration file. 36 | func EnvFile() (string, error) { 37 | if file := os.Getenv("GOENV"); file != "" { 38 | if file == "off" { 39 | return "", fmt.Errorf("GOENV=off") 40 | } 41 | return file, nil 42 | } 43 | dir, err := os.UserConfigDir() 44 | if err != nil { 45 | return "", err 46 | } 47 | if dir == "" { 48 | return "", fmt.Errorf("missing user-config dir") 49 | } 50 | return filepath.Join(dir, "go/env"), nil 51 | } 52 | 53 | // initEnvCache init environment cahce 54 | func initEnvCache() { 55 | envCache.m = make(map[string]string) 56 | file, _ := EnvFile() 57 | if file == "" { 58 | return 59 | } 60 | data, err := ioutil.ReadFile(file) 61 | if err != nil { 62 | return 63 | } 64 | 65 | for len(data) > 0 { 66 | // Get next line. 67 | line := data 68 | i := bytes.IndexByte(data, '\n') 69 | if i >= 0 { 70 | line, data = line[:i], data[i+1:] 71 | } else { 72 | data = nil 73 | } 74 | 75 | i = bytes.IndexByte(line, '=') 76 | if i < 0 || line[0] < 'A' || 'Z' < line[0] { 77 | // Line is missing = (or empty) or a comment or not a valid env name. Ignore. 78 | // (This should not happen, since the file should be maintained almost 79 | // exclusively by "go env -w", but better to silently ignore than to make 80 | // the go command unusable just because somehow the env file has 81 | // gotten corrupted.) 82 | continue 83 | } 84 | key, val := line[:i], line[i+1:] 85 | envCache.m[string(key)] = string(val) 86 | } 87 | } 88 | 89 | // Getenv gets the value from env or configuration. 90 | func Getenv(key string) string { 91 | val := os.Getenv(key) 92 | if val != "" { 93 | return val 94 | } 95 | envCache.once.Do(initEnvCache) 96 | return envCache.m[key] 97 | } 98 | -------------------------------------------------------------------------------- /tools/waterdrop/utils/helper.go: -------------------------------------------------------------------------------- 1 | /* 2 | * 3 | * Copyright 2020 waterdrop authors. 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | * 17 | */ 18 | 19 | package utils 20 | 21 | import ( 22 | "fmt" 23 | "go/build" 24 | "os" 25 | "os/exec" 26 | "path/filepath" 27 | "runtime" 28 | "strings" 29 | 30 | "github.com/urfave/cli/v2" 31 | ) 32 | 33 | // ExecuteGoGet execute go get command 34 | func ExecuteGoGet(address string) error { 35 | args := strings.Split(address, " ") 36 | cmd := exec.Command(args[0], args[1:]...) 37 | cmd.Env = os.Environ() 38 | cmd.Stdout = os.Stdout 39 | cmd.Stderr = os.Stderr 40 | return cmd.Run() 41 | } 42 | 43 | // RunTool run tool commands 44 | func RunTool(ctx *cli.Context, dir string, args []string) (err error) { 45 | cmd := toolPath(ctx.Command.Name) 46 | toolCmd := &exec.Cmd{ 47 | Path: cmd, 48 | Args: append([]string{cmd}, args...), 49 | Dir: dir, 50 | Stdin: os.Stdin, 51 | Stdout: os.Stdout, 52 | Stderr: os.Stderr, 53 | Env: os.Environ(), 54 | } 55 | if filepath.Base(cmd) == cmd { 56 | var lp string 57 | if lp, err = exec.LookPath(cmd); err == nil { 58 | toolCmd.Path = lp 59 | } 60 | } 61 | if err = toolCmd.Run(); err != nil { 62 | if e, ok := err.(*exec.ExitError); !ok || !e.Exited() { 63 | fmt.Fprintf(os.Stderr, "运行 %s 出错: %v\n", ctx.Command.Name, err) 64 | } 65 | } 66 | return 67 | } 68 | 69 | // Gopath return go path directory 70 | func Gopath() (gp string) { 71 | gopaths := strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator)) 72 | 73 | if len(gopaths) == 1 && gopaths[0] != "" { 74 | return gopaths[0] 75 | } 76 | pwd, err := os.Getwd() 77 | if err != nil { 78 | return 79 | } 80 | abspwd, err := filepath.Abs(pwd) 81 | if err != nil { 82 | return 83 | } 84 | for _, gopath := range gopaths { 85 | if gopath == "" { 86 | continue 87 | } 88 | absgp, err := filepath.Abs(gopath) 89 | if err != nil { 90 | return 91 | } 92 | if strings.HasPrefix(abspwd, absgp) { 93 | return absgp 94 | } 95 | } 96 | return build.Default.GOPATH 97 | } 98 | 99 | // toolPath return tool path 100 | func toolPath(toolName string) string { 101 | gobin := Getenv("GOBIN") 102 | if runtime.GOOS == "windows" { 103 | toolName += ".exe" 104 | } 105 | if gobin != "" { 106 | return filepath.Join(gobin, toolName) 107 | } 108 | return filepath.Join(Gopath(), "bin", toolName) 109 | } 110 | --------------------------------------------------------------------------------